diff --git a/pom.xml b/pom.xml
index e8c2b94..0e8c3d2 100644
--- a/pom.xml
+++ b/pom.xml
@@ -6,7 +6,7 @@
org.springframework.boot
spring-boot-starter-parent
- 3.1.3
+ 3.2.0
com.apolloconfig.apollo.ai
@@ -16,10 +16,13 @@
a smart qa bot
17
- 0.16.0
- 32.1.2-jre
+ 0.18.2
+ 32.1.3-jre
0.64.8
- 2.3.0
+ 2.3.3
+
+ 4.1.100.Final
+ 1.4.0
@@ -44,6 +47,16 @@
milvus-sdk-java
${milvus.version}
+
+ io.netty
+ netty-codec-http2
+ ${netty.codec.http2.version}
+
+
+ org.crac
+ crac
+ ${crac.version}
+
@@ -82,6 +95,10 @@
+
+ org.crac
+ crac
+
org.springframework.boot
diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java
index 24d685e..43be515 100644
--- a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java
+++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusClientFactory.java
@@ -4,8 +4,10 @@
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import java.util.Map;
+import org.crac.Context;
+import org.crac.Resource;
-public class MilvusClientFactory {
+public class MilvusClientFactory implements Resource {
private static final MilvusClientFactory INSTANCE = new MilvusClientFactory();
private static final Map clients = Maps.newConcurrentMap();
@@ -43,4 +45,14 @@ private MilvusServiceClient createClient(String host, int port) {
private MilvusServiceClient createCloudClient(String uri, String token) {
return new MilvusServiceClient(ConnectParam.newBuilder().withUri(uri).withToken(token).build());
}
+
+ @Override
+ public void beforeCheckpoint(Context extends Resource> context) throws Exception {
+ clients.clear();
+ }
+
+ @Override
+ public void afterRestore(Context extends Resource> context) throws Exception {
+
+ }
}
diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java
index 6becc6e..e5c541b 100644
--- a/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java
+++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/milvus/MilvusService.java
@@ -34,20 +34,30 @@
import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;
+import org.crac.Context;
+import org.crac.Resource;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
@Profile("milvus")
@Service
-class MilvusService implements VectorDBService {
+class MilvusService implements VectorDBService, Resource {
- private final MilvusServiceClient milvusServiceClient;
+ private static final Logger LOGGER = LoggerFactory.getLogger(MilvusService.class);
+
+ private MilvusServiceClient milvusServiceClient;
private final MilvusConfig milvusConfig;
private final List dummyEmbeddings = Lists.newArrayList();
public MilvusService(MilvusConfig milvusConfig) {
this.milvusConfig = milvusConfig;
+ this.init();
+ }
+
+ private void init() {
if (milvusConfig.isUseZillzCloud()) {
this.milvusServiceClient = MilvusClientFactory.getCloudClient(
milvusConfig.getZillizCloudUri(),
@@ -413,4 +423,16 @@ private void ensureFileCollection() {
);
}
+ @Override
+ public void beforeCheckpoint(Context extends Resource> context) throws Exception {
+ LOGGER.info("beforeCheckpoint");
+ this.milvusServiceClient = null;
+ }
+
+ @Override
+ public void afterRestore(Context extends Resource> context) throws Exception {
+ LOGGER.info("afterRestore");
+ this.init();
+ LOGGER.info("afterRestore done");
+ }
}
diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java
index f02028c..c8e2344 100644
--- a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java
+++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiService.java
@@ -10,19 +10,25 @@
import com.theokanning.openai.embedding.EmbeddingRequest;
import io.reactivex.Flowable;
import java.util.List;
+import org.crac.Context;
+import org.crac.Resource;
import org.springframework.context.annotation.Profile;
import org.springframework.stereotype.Component;
@Profile("openai")
@Component
-class OpenAiService implements AiService {
+class OpenAiService implements AiService, Resource {
private static final String DEFAULT_MODEL = "gpt-3.5-turbo";
private static final String DEFAULT_EMBEDDING_MODEL = "text-embedding-ada-002";
- private final com.theokanning.openai.service.OpenAiService service;
+ private com.theokanning.openai.service.OpenAiService service;
public OpenAiService() {
+ init();
+ }
+
+ private void init() {
service = OpenAiServiceFactory.getService(System.getenv("OPENAI_API_KEY"));
}
@@ -60,4 +66,14 @@ public List getEmbeddings(List chunks) {
return service.createEmbeddings(embeddingRequest).getData();
}
+
+ @Override
+ public void beforeCheckpoint(Context extends Resource> context) throws Exception {
+ this.service = null;
+ }
+
+ @Override
+ public void afterRestore(Context extends Resource> context) throws Exception {
+ this.init();
+ }
}
diff --git a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java
index 1d67662..770b979 100644
--- a/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java
+++ b/src/main/java/com/apolloconfig/apollo/ai/qabot/openai/OpenAiServiceFactory.java
@@ -23,9 +23,11 @@
import okhttp3.Authenticator;
import okhttp3.Credentials;
import okhttp3.OkHttpClient;
+import org.crac.Context;
+import org.crac.Resource;
import retrofit2.Retrofit;
-public class OpenAiServiceFactory {
+public class OpenAiServiceFactory implements Resource {
private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(60);
@@ -85,6 +87,16 @@ private OkHttpClient client(String apiKey) {
.build();
}
+ @Override
+ public void beforeCheckpoint(Context extends Resource> context) throws Exception {
+ SERVICES.clear();
+ }
+
+ @Override
+ public void afterRestore(Context extends Resource> context) throws Exception {
+
+ }
+
private static class DelegatingSocketFactory extends SocketFactory {
private final SocketFactory delegate;
diff --git a/src/main/scripts/startup.sh b/src/main/scripts/startup.sh
index 77f1b54..00a2130 100644
--- a/src/main/scripts/startup.sh
+++ b/src/main/scripts/startup.sh
@@ -19,13 +19,27 @@ SERVICE_NAME=qa-bot
LOG_DIR=/opt/logs
## Adjust server port if necessary
SERVER_PORT=${SERVER_PORT:=9090}
+## Adjust crac files dir if necessary
+CRAC_FILES_DIR=/opt/crac
## Create log directory if not existed because JDK 8+ won't do that
mkdir -p $LOG_DIR
+mkdir -p $CRAC_FILES_DIR
+
## Adjust memory settings if necessary
#export JAVA_OPTS="-Xms2560m -Xmx2560m -Xss256k -XX:MetaspaceSize=128m -XX:MaxMetaspaceSize=384m -XX:NewSize=1536m -XX:MaxNewSize=1536m -XX:SurvivorRatio=8"
+# Check for 'checkpoint' argument
+if [ "$1" = "checkpoint" ]; then
+ export JAVA_OPTS="$JAVA_OPTS -Dspring.context.checkpoint=onRefresh -XX:CRaCCheckpointTo=$CRAC_FILES_DIR"
+fi
+
+# Check for 'restore' argument
+if [ "$1" = "restore" ]; then
+ export JAVA_OPTS="$JAVA_OPTS -XX:CRaCRestoreFrom=$CRAC_FILES_DIR"
+fi
+
## Only uncomment the following when you are using server jvm
export JAVA_OPTS="$JAVA_OPTS -server -XX:-ReduceInitialCardMarks"