feat(ai): Adding Lucene & Embedding-Based Search Operators to Apache GeaFlow (incubating) for Lightweight Context Memory

geaflow-ai/pom.xml

@@ -0,0 +1,127 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ 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.
+  -->
+
+<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xmlns="http://maven.apache.org/POM/4.0.0"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.apache.geaflow</groupId>
+        <artifactId>geaflow-parent</artifactId>
+        <version>0.8.0-SNAPSHOT</version>
+    </parent>
+
+    <artifactId>geaflow-ai</artifactId>
+
+    <properties>
+        <solon-mcp.version>3.5.1</solon-mcp.version>
+        <junit.version>5.10.1</junit.version>
+        <slf4j.version>1.7.15</slf4j.version>
+        <log4j.version>1.2.17</log4j.version>
+        <log4j.slf4j.version>2.17.1</log4j.slf4j.version>
+        <lmax.disrupter.veresion>3.4.4</lmax.disrupter.veresion>
+        <lucene.version>8.11.2</lucene.version>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+            <version>${slf4j.version}</version>
+        </dependency>
+
+        <!-- SLF4J to Log4j 1.2 Binding -->
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-log4j12</artifactId>
+            <version>${slf4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>log4j</groupId>
+            <artifactId>log4j</artifactId>
+            <version>${log4j.version}</version>
+        </dependency>
+
+        <!-- SLF4J to Log4j 2 Binding -->
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-slf4j-impl</artifactId>
+            <version>${log4j.slf4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-core</artifactId>
+            <version>${log4j.slf4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-api</artifactId>
+            <version>${log4j.slf4j.version}</version>
+        </dependency>
+
+        <!-- Lucene Core -->
+        <dependency>
+            <groupId>org.apache.lucene</groupId>
+            <artifactId>lucene-core</artifactId>
+            <version>${lucene.version}</version>
+        </dependency>
+
+        <!-- Lucene Query Parser -->
+        <dependency>
+            <groupId>org.apache.lucene</groupId>
+            <artifactId>lucene-queryparser</artifactId>
+            <version>${lucene.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.lmax</groupId>
+            <artifactId>disruptor</artifactId>
+            <version>${lmax.disrupter.veresion}</version>
+        </dependency>
+
+      <dependency>
+        <groupId>org.apache.geaflow</groupId>
+        <artifactId>geaflow-api</artifactId>
+        <version>${project.version}</version>
+      </dependency>
+
+        <!-- test -->
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <version>${junit.version}</version>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <repositories>
+        <repository>
+            <id>sonatype</id>
+            <url>https://s01.oss.sonatype.org/content/repositories/snapshots</url>
+            <releases>
+                <enabled>false</enabled>
+            </releases>
+            <snapshots>
+                <enabled>true</enabled>
+            </snapshots>
+        </repository>
+    </repositories>
+
+</project>

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

@@ -0,0 +1,110 @@
+/*
+ * 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;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.geaflow.ai.graph.GraphAccessor;
+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.SearchOperator;
+import org.apache.geaflow.ai.operator.SessionOperator;
+import org.apache.geaflow.ai.search.VectorSearch;
+import org.apache.geaflow.ai.session.SessionManagement;
+import org.apache.geaflow.ai.subgraph.SubGraph;
+import org.apache.geaflow.ai.verbalization.Context;
+import org.apache.geaflow.ai.verbalization.VerbalizationFunction;
+
+public class GraphMemoryServer {
+
+    private final SessionManagement sessionManagement = SessionManagement.INSTANCE;
+    private final List<GraphAccessor> graphAccessors = new ArrayList<>();
+    private final List<IndexStore> indexStores = new ArrayList<>();
+
+    public void addGraphAccessor(GraphAccessor graph) {
+        if (graph != null) {
+            graphAccessors.add(graph);
+        }
+    }
+
+    public void addIndexStore(IndexStore indexStore) {
+        if (indexStore != null) {
+            indexStores.add(indexStore);
+        }
+    }
+
+    public String createSession() {
+        String sessionId = sessionManagement.createSession();
+        if (sessionId == null) {
+            throw new RuntimeException("Cannot create new session");
+        }
+        return sessionId;
+    }
+
+    public String search(VectorSearch search) {
+        String sessionId = search.getSessionId();
+        if (sessionId == null || sessionId.isEmpty()) {
+            throw new RuntimeException("Session id is empty");
+        }
+        if (!sessionManagement.sessionExists(sessionId)) {
+            sessionManagement.createSession(sessionId);
+        }
+
+        for (IndexStore indexStore : indexStores) {
+            if (indexStore instanceof EntityAttributeIndexStore) {
+                SessionOperator searchOperator = new SessionOperator(graphAccessors.get(0), indexStore);
+                applySearch(sessionId, searchOperator, search);
+            }
+            if (indexStore instanceof EmbeddingIndexStore) {
+                EmbeddingOperator embeddingOperator = new EmbeddingOperator(graphAccessors.get(0), indexStore);
+                applySearch(sessionId, embeddingOperator, search);
+            }
+        }
+        return sessionId;
+    }
+
+    private void applySearch(String sessionId, SearchOperator operator, VectorSearch search) {
+        SessionManagement manager = SessionManagement.INSTANCE;
+        if (!manager.sessionExists(sessionId)) {
+            return;
+        }
+        List<SubGraph> result = operator.apply(manager.getSubGraph(sessionId), search);
+        manager.setSubGraph(sessionId, result);
+    }
+
+    public Context verbalize(String sessionId, VerbalizationFunction verbalizationFunction) {
+        List<SubGraph> subGraphList = sessionManagement.getSubGraph(sessionId);
+        List<String> subGraphStringList = new ArrayList<>(subGraphList.size());
+        for (SubGraph subGraph : subGraphList) {
+            subGraphStringList.add(verbalizationFunction.verbalize(subGraph));
+        }
+        subGraphStringList = subGraphStringList.stream().sorted().collect(Collectors.toList());
+        StringBuilder stringBuilder = new StringBuilder();
+        for (String subGraph : subGraphStringList) {
+            stringBuilder.append(subGraph).append("\n");
+        }
+        stringBuilder.append(verbalizationFunction.verbalizeGraphSchema());
+        return new Context(stringBuilder.toString());
+    }
+
+}

geaflow-ai/src/main/java/org/apache/geaflow/ai/common/ErrorCode.java

@@ -0,0 +1,30 @@
+/*
+ * 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.common;
+
+public class ErrorCode {
+
+    public static final int SUCCESS = 0;
+    public static final int GRAPH_ENTITY_GROUP_NOT_EXISTS = 100001;
+    public static final int GRAPH_ENTITY_GROUP_NOT_MATCH = 100002;
+    public static final int GRAPH_ENTITY_GROUP_INSERT_FAILED = 100003;
+    public static final int GRAPH_ENTITY_GROUP_UPDATE_FAILED = 100004;
+    public static final int GRAPH_ENTITY_GROUP_REMOVE_FAILED = 100005;
+}

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

@@ -0,0 +1,46 @@
+/*
+ * 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.common.config;
+
+public class Constants {
+
+    public static String MODEL_CONTEXT_ROLE_USER = "user";
+    public static String PREFIX_V = "V";
+    public static String PREFIX_E = "E";
+    public static String PREFIX_GRAPH = "GRAPH";
+    public static String PREFIX_TMP_SESSION = "TmpSession-";
+
+    public static int HTTP_CALL_TIMEOUT_SECONDS = 300;
+    public static int HTTP_CONNECT_TIMEOUT_SECONDS = 300;
+    public static int HTTP_READ_TIMEOUT_SECONDS = 300;
+    public static int HTTP_WRITE_TIMEOUT_SECONDS = 300;
+
+    public static int MODEL_CLIENT_RETRY_TIMES = 10;
+    public static int MODEL_CLIENT_RETRY_INTERVAL_MS = 3000;
+
+    public static int EMBEDDING_INDEX_STORE_BATCH_SIZE = 32;
+    public static int EMBEDDING_INDEX_STORE_REPORT_SIZE = 100;
+    public static int EMBEDDING_INDEX_STORE_FLUSH_WRITE_SIZE = 1024;
+    public static int EMBEDDING_INDEX_STORE_SPLIT_TEXT_CHUNK_SIZE = 128;
+
+    public static double EMBEDDING_OPERATE_DEFAULT_THRESHOLD = 0.5;
+    public static int EMBEDDING_OPERATE_DEFAULT_TOPN = 50;
+    public static int GRAPH_SEARCH_STORE_DEFAULT_TOPN = 30;
+}

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

@@ -0,0 +1,23 @@
+/*
+ * 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.common.config;
+
+public class GraphMemoryConfigKeys {
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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.common.model;
+
+public abstract class AbstractModelService {
+
+    public ModelConfig getModelConfig() {
+        return modelConfig;
+    }
+
+    public void setModelConfig(ModelConfig modelConfig) {
+        this.modelConfig = modelConfig;
+    }
+
+    private ModelConfig modelConfig;
+
+    public AbstractModelService(ModelConfig modelConfig) {
+        this.modelConfig = modelConfig;
+    }
+}