Skip to content

Commit 72be83b

Browse files
LeomrlinLeomrlin
authored
feat(ai): support graph consolidate algorithm (#729)
* init dcp code * add lucene search * add prompt formatter * add test case * handle ldbc id conflict * support llm * support embedding index store * add embedding op * refine test case * delete test data * add MockChatRobot * fix checkstyle * fix pom * fix finishReason * fix ci tests * fix ci tests * fix comments * fix codestyle * support mutable graph * Add ConsolidateFunction * Add geaflow memory server * Add GeaFlowMemoryClientCLI * fix comment * refine code * replay commit * add consolidate server
1 parent ee4a0c5 commit 72be83b

23 files changed

Lines changed: 2640 additions & 17 deletions

‎geaflow-ai/pom.xml‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,14 @@
3131
<artifactId>geaflow-ai</artifactId>
3232

3333
<properties>
34-
<solon-mcp.version>3.5.1</solon-mcp.version>
34+
<solon.version>3.5.1</solon.version>
3535
<junit.version>5.10.1</junit.version>
3636
<slf4j.version>1.7.15</slf4j.version>
3737
<log4j.version>1.2.17</log4j.version>
3838
<log4j.slf4j.version>2.17.1</log4j.slf4j.version>
3939
<lmax.disrupter.veresion>3.4.4</lmax.disrupter.veresion>
4040
<lucene.version>8.11.2</lucene.version>
41+
<gson.version>2.2.4</gson.version>
4142
</properties>
4243

4344
<dependencies>
@@ -96,6 +97,25 @@
9697
<version>${lmax.disrupter.veresion}</version>
9798
</dependency>
9899

100+
<dependency>
101+
<groupId>org.noear</groupId>
102+
<artifactId>solon-web</artifactId>
103+
<version>${solon.version}</version>
104+
</dependency>
105+
106+
<dependency>
107+
<groupId>com.google.code.gson</groupId>
108+
<artifactId>gson</artifactId>
109+
<version>${gson.version}</version>
110+
</dependency>
111+
112+
<dependency>
113+
<groupId>org.noear</groupId>
114+
<artifactId>solon-test</artifactId>
115+
<version>${solon.version}</version>
116+
<scope>test</scope>
117+
</dependency>
118+
99119
<dependency>
100120
<groupId>org.apache.geaflow</groupId>
101121
<artifactId>geaflow-api</artifactId>
@@ -124,4 +144,30 @@
124144
</repository>
125145
</repositories>
126146

147+
<build>
148+
<plugins>
149+
<plugin>
150+
<artifactId>maven-assembly-plugin</artifactId>
151+
<configuration>
152+
<archive>
153+
<manifest>
154+
<mainClass>org.apache.geaflow.ai.client.GeaFlowMemoryClientCLI</mainClass>
155+
</manifest>
156+
</archive>
157+
<descriptorRefs>
158+
<descriptorRef>jar-with-dependencies</descriptorRef>
159+
</descriptorRefs>
160+
</configuration>
161+
<executions>
162+
<execution>
163+
<id>make-assembly</id>
164+
<phase>package</phase>
165+
<goals>
166+
<goal>single</goal>
167+
</goals>
168+
</execution>
169+
</executions>
170+
</plugin>
171+
</plugins>
172+
</build>
127173
</project>

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

Lines changed: 232 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,232 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.geaflow.ai;
21+
22+
import java.util.ArrayList;
23+
import java.util.HashMap;
24+
import java.util.List;
25+
import java.util.Map;
26+
import org.apache.geaflow.ai.common.util.SeDeUtil;
27+
import org.apache.geaflow.ai.graph.*;
28+
import org.apache.geaflow.ai.graph.io.*;
29+
import org.apache.geaflow.ai.index.EntityAttributeIndexStore;
30+
import org.apache.geaflow.ai.index.vector.KeywordVector;
31+
import org.apache.geaflow.ai.search.VectorSearch;
32+
import org.apache.geaflow.ai.service.ServerMemoryCache;
33+
import org.apache.geaflow.ai.verbalization.Context;
34+
import org.apache.geaflow.ai.verbalization.SubgraphSemanticPromptFunction;
35+
import org.noear.solon.Solon;
36+
import org.noear.solon.annotation.*;
37+
import org.slf4j.Logger;
38+
import org.slf4j.LoggerFactory;
39+
40+
@Controller
41+
public class GeaFlowMemoryServer {
42+
43+
private static final Logger LOGGER = LoggerFactory.getLogger(GeaFlowMemoryServer.class);
44+
45+
private static final String SERVER_NAME = "geaflow-memory-server";
46+
private static final int DEFAULT_PORT = 8080;
47+
48+
private static final ServerMemoryCache CACHE = new ServerMemoryCache();
49+
50+
public static void main(String[] args) {
51+
System.setProperty("solon.app.name", SERVER_NAME);
52+
Solon.start(GeaFlowMemoryServer.class, args, app -> {
53+
app.cfg().loadAdd("application.yml");
54+
int port = app.cfg().getInt("server.port", DEFAULT_PORT);
55+
LOGGER.info("Starting {} on port {}", SERVER_NAME, port);
56+
app.get("/", ctx -> {
57+
ctx.output("GeaFlow AI Server is running...");
58+
});
59+
app.get("/health", ctx -> {
60+
ctx.output("{\"status\":\"UP\",\"service\":\"" + SERVER_NAME + "\"}");
61+
});
62+
});
63+
}
64+
65+
@Get
66+
@Mapping("/api/test")
67+
public String test() {
68+
return "GeaFlow Memory Server is working!";
69+
}
70+
71+
@Post
72+
@Mapping("/graph/create")
73+
public String createGraph(@Body String input) {
74+
GraphSchema graphSchema = SeDeUtil.deserializeGraphSchema(input);
75+
String graphName = graphSchema.getName();
76+
if (graphName == null || CACHE.getGraphByName(graphName) != null) {
77+
throw new RuntimeException("Cannot create graph name: " + graphName);
78+
}
79+
Map<String, EntityGroup> entities = new HashMap<>();
80+
for (VertexSchema vertexSchema : graphSchema.getVertexSchemaList()) {
81+
entities.put(vertexSchema.getName(), new VertexGroup(vertexSchema, new ArrayList<>()));
82+
}
83+
for (EdgeSchema edgeSchema : graphSchema.getEdgeSchemaList()) {
84+
entities.put(edgeSchema.getName(), new EdgeGroup(edgeSchema, new ArrayList<>()));
85+
}
86+
MemoryGraph graph = new MemoryGraph(graphSchema, entities);
87+
CACHE.putGraph(graph);
88+
LocalMemoryGraphAccessor graphAccessor = new LocalMemoryGraphAccessor(graph);
89+
LOGGER.info("Success to init empty graph.");
90+
91+
EntityAttributeIndexStore indexStore = new EntityAttributeIndexStore();
92+
indexStore.initStore(new SubgraphSemanticPromptFunction(graphAccessor));
93+
LOGGER.info("Success to init EntityAttributeIndexStore.");
94+
95+
GraphMemoryServer server = new GraphMemoryServer();
96+
server.addGraphAccessor(graphAccessor);
97+
server.addIndexStore(indexStore);
98+
LOGGER.info("Success to init GraphMemoryServer.");
99+
CACHE.putServer(server);
100+
101+
LOGGER.info("Success to init graph. SCHEMA: {}", graphSchema);
102+
return "createGraph has been called, graphName: " + graphName;
103+
}
104+
105+
@Post
106+
@Mapping("/graph/addEntitySchema")
107+
public String addSchema(@Param("graphName") String graphName,
108+
@Body String input) {
109+
Graph graph = CACHE.getGraphByName(graphName);
110+
if (graph == null) {
111+
throw new RuntimeException("Graph not exist.");
112+
}
113+
if (!(graph instanceof MemoryGraph)) {
114+
throw new RuntimeException("Graph cannot modify.");
115+
}
116+
MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph);
117+
Schema schema = SeDeUtil.deserializeEntitySchema(input);
118+
String schemaName = schema.getName();
119+
if (schema instanceof VertexSchema) {
120+
memoryMutableGraph.addVertexSchema((VertexSchema) schema);
121+
} else if (schema instanceof EdgeSchema) {
122+
memoryMutableGraph.addEdgeSchema((EdgeSchema) schema);
123+
} else {
124+
throw new RuntimeException("Cannt add schema: " + input);
125+
}
126+
return "addSchema has been called, schemaName: " + schemaName;
127+
}
128+
129+
@Post
130+
@Mapping("/graph/getGraphSchema")
131+
public String getSchema(@Param("graphName") String graphName) {
132+
Graph graph = CACHE.getGraphByName(graphName);
133+
if (graph == null) {
134+
throw new RuntimeException("Graph not exist.");
135+
}
136+
if (!(graph instanceof MemoryGraph)) {
137+
throw new RuntimeException("Graph cannot modify.");
138+
}
139+
return SeDeUtil.serializeGraphSchema(graph.getGraphSchema());
140+
}
141+
142+
@Post
143+
@Mapping("/graph/insertEntity")
144+
public String addEntity(@Param("graphName") String graphName,
145+
@Body String input) {
146+
Graph graph = CACHE.getGraphByName(graphName);
147+
if (graph == null) {
148+
throw new RuntimeException("Graph not exist.");
149+
}
150+
if (!(graph instanceof MemoryGraph)) {
151+
throw new RuntimeException("Graph cannot modify.");
152+
}
153+
MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph);
154+
List<GraphEntity> graphEntities = SeDeUtil.deserializeEntities(input);
155+
156+
for (GraphEntity entity : graphEntities) {
157+
if (entity instanceof GraphVertex) {
158+
memoryMutableGraph.addVertex(((GraphVertex) entity).getVertex());
159+
} else {
160+
memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge());
161+
}
162+
}
163+
CACHE.getConsolidateServer().executeConsolidateTask(
164+
CACHE.getServerByName(graphName).getGraphAccessors().get(0), memoryMutableGraph);
165+
return "Success to add entities, num: " + graphEntities.size();
166+
}
167+
168+
@Post
169+
@Mapping("/graph/delEntity")
170+
public String deleteEntity(@Param("graphName") String graphName,
171+
@Body String input) {
172+
Graph graph = CACHE.getGraphByName(graphName);
173+
if (graph == null) {
174+
throw new RuntimeException("Graph not exist.");
175+
}
176+
if (!(graph instanceof MemoryGraph)) {
177+
throw new RuntimeException("Graph cannot modify.");
178+
}
179+
MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph);
180+
List<GraphEntity> graphEntities = SeDeUtil.deserializeEntities(input);
181+
for (GraphEntity entity : graphEntities) {
182+
if (entity instanceof GraphVertex) {
183+
memoryMutableGraph.removeVertex(entity.getLabel(),
184+
((GraphVertex) entity).getVertex().getId());
185+
} else {
186+
memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge());
187+
}
188+
}
189+
return "Success to remove entities, num: " + graphEntities.size();
190+
}
191+
192+
@Post
193+
@Mapping("/query/context")
194+
public String createContext(@Param("graphName") String graphName) {
195+
GraphMemoryServer server = CACHE.getServerByName(graphName);
196+
if (server == null) {
197+
throw new RuntimeException("Server not exist.");
198+
}
199+
String sessionId = server.createSession();
200+
CACHE.putSession(server, sessionId);
201+
return sessionId;
202+
}
203+
204+
@Post
205+
@Mapping("/query/exec")
206+
public String execQuery(@Param("sessionId") String sessionId,
207+
@Body String query) {
208+
String graphName = CACHE.getGraphNameBySession(sessionId);
209+
if (graphName == null) {
210+
throw new RuntimeException("Graph not exist.");
211+
}
212+
GraphMemoryServer server = CACHE.getServerByName(graphName);
213+
VectorSearch search = new VectorSearch(null, sessionId);
214+
search.addVector(new KeywordVector(query));
215+
server.search(search);
216+
Context context = server.verbalize(sessionId,
217+
new SubgraphSemanticPromptFunction(server.getGraphAccessors().get(0)));
218+
return context.toString();
219+
}
220+
221+
@Post
222+
@Mapping("/query/result")
223+
public String getResult(@Param("sessionId") String sessionId) {
224+
String graphName = CACHE.getGraphNameBySession(sessionId);
225+
if (graphName == null) {
226+
throw new RuntimeException("Graph not exist.");
227+
}
228+
GraphMemoryServer server = CACHE.getServerByName(graphName);
229+
List<GraphEntity> result = server.getSessionEntities(sessionId);
230+
return result.toString();
231+
}
232+
}

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

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,12 @@
2020
package org.apache.geaflow.ai;
2121

2222
import java.util.ArrayList;
23+
import java.util.HashSet;
2324
import java.util.List;
25+
import java.util.Set;
2426
import java.util.stream.Collectors;
2527
import org.apache.geaflow.ai.graph.GraphAccessor;
28+
import org.apache.geaflow.ai.graph.GraphEntity;
2629
import org.apache.geaflow.ai.index.EmbeddingIndexStore;
2730
import org.apache.geaflow.ai.index.EntityAttributeIndexStore;
2831
import org.apache.geaflow.ai.index.IndexStore;
@@ -37,7 +40,7 @@
3740

3841
public class GraphMemoryServer {
3942

40-
private final SessionManagement sessionManagement = SessionManagement.INSTANCE;
43+
private final SessionManagement sessionManagement = new SessionManagement();
4144
private final List<GraphAccessor> graphAccessors = new ArrayList<>();
4245
private final List<IndexStore> indexStores = new ArrayList<>();
4346

@@ -47,12 +50,20 @@ public void addGraphAccessor(GraphAccessor graph) {
4750
}
4851
}
4952

53+
public List<GraphAccessor> getGraphAccessors() {
54+
return graphAccessors;
55+
}
56+
5057
public void addIndexStore(IndexStore indexStore) {
5158
if (indexStore != null) {
5259
indexStores.add(indexStore);
5360
}
5461
}
5562

63+
public List<IndexStore> getIndexStores() {
64+
return indexStores;
65+
}
66+
5667
public String createSession() {
5768
String sessionId = sessionManagement.createSession();
5869
if (sessionId == null) {
@@ -84,7 +95,7 @@ public String search(VectorSearch search) {
8495
}
8596

8697
private void applySearch(String sessionId, SearchOperator operator, VectorSearch search) {
87-
SessionManagement manager = SessionManagement.INSTANCE;
98+
SessionManagement manager = sessionManagement;
8899
if (!manager.sessionExists(sessionId)) {
89100
return;
90101
}
@@ -107,4 +118,13 @@ public Context verbalize(String sessionId, VerbalizationFunction verbalizationFu
107118
return new Context(stringBuilder.toString());
108119
}
109120

121+
public List<GraphEntity> getSessionEntities(String sessionId) {
122+
List<SubGraph> subGraphList = sessionManagement.getSubGraph(sessionId);
123+
Set<GraphEntity> entitySet = new HashSet<>();
124+
for (SubGraph subGraph : subGraphList) {
125+
entitySet.addAll(subGraph.getGraphEntityList());
126+
}
127+
return new ArrayList<>(entitySet);
128+
}
129+
110130
}

0 commit comments

Comments
 (0)