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;
+    }
+
+    public 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/model/ChatRobot.java

@@ -0,0 +1,84 @@
+/*
+ * 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;
+
+import com.google.gson.Gson;
+import java.util.List;
+
+public class ChatRobot {
+
+    private ModelInfo modelInfo;
+
+    public ChatRobot() {
+        this.modelInfo = new ModelInfo();
+    }
+
+    public ChatRobot(String model) {
+        this.modelInfo = new ModelInfo();
+        this.modelInfo.setModel(model);
+    }
+
+    public String singleSentence(String sentence) {
+        OfflineModelDirect model = new OfflineModelDirect();
+        ModelContext context = ModelContext.emptyContext();
+        context.setModelInfo(modelInfo);
+        context.userSay(sentence);
+        return model.chat(context);
+    }
+
+    public String embedding(String... inputs) {
+        OfflineModelDirect model = new OfflineModelDirect();
+        ModelEmbedding context = ModelEmbedding.embedding(modelInfo, inputs);
+        List<EmbeddingResult> embeddingResults = model.embedding(context);
+        Gson gson = new Gson();
+        StringBuilder builder = new StringBuilder();
+        for (EmbeddingResult result : embeddingResults) {
+            builder.append("\n");
+            String json = gson.toJson(result);
+            builder.append(json);
+        }
+        return builder.toString();
+    }
+
+    public EmbeddingResult embeddingSingle(String input) {
+        OfflineModelDirect model = new OfflineModelDirect();
+        ModelEmbedding context = ModelEmbedding.embedding(modelInfo, input);
+        List<EmbeddingResult> embeddingResults = model.embedding(context);
+        return embeddingResults.get(0);
+    }
+
+    public ModelInfo getModelInfo() {
+        return modelInfo;
+    }
+
+    public void setModelInfo(ModelInfo modelInfo) {
+        this.modelInfo = modelInfo;
+    }
+
+    public static class EmbeddingResult {
+        public String input;
+        public double[] embedding;
+
+        public EmbeddingResult(String input, double[] embedding) {
+            this.input = input;
+            this.embedding = embedding;
+        }
+    }
+}

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

@@ -0,0 +1,47 @@
+/*
+ * 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;
+
+import com.google.gson.annotations.SerializedName;
+import java.util.List;
+
+public class EmbeddingResponse {
+
+    public List<EmbeddingVector> data;
+    public String model;
+    public String object;
+    public Response.Usage usage;
+
+    public static class EmbeddingVector {
+        public int index;
+        public double[] embedding;
+    }
+
+    public static class Usage {
+        @SerializedName("completion_tokens")
+        public int completionTokens;
+        @SerializedName("prompt_tokens")
+        public int promptTokens;
+        @SerializedName("total_tokens")
+        public int totalTokens;
+
+    }
+
+}