feature: store and load conversations

Author: kezhenxu94

README.md

@@ -26,18 +26,27 @@ implementation 'io.github.kezhenxu94:chatgpt-java-sdk:${chatgpt-java.version}'
 ```java
 public class ChatGPTTest {
     public static void main(String[] args) throws IOException, InterruptedException {
-        final var chatGPT = ChatGPT.builder().apiKey("").build(); // Setting API Key via environment variable (CHATGPT_API_KEY) is also supported.
+        final var chatGPT = ChatGPT
+                .builder()
+                .dataPath(Files.createTempDirectory("chatgpt")) // Persist the chat history to a data path
+                .build();
+
+        // Start a new conversation
         final var conversation = chatGPT.newConversation();
         System.out.println(conversation.ask("What's your name?").content());
         // Output: I'm an AI language model developed by OpenAI, and I don't have a name. What can I help you with today?
         System.out.println(conversation.ask("What did I ask you?").content());
         // Output: You asked for my name.
+        conversation.save(); // Save the history manually, conversations are saved on shutdown by default.
 
         final var conversation2 = chatGPT.newConversation("You are a software engineer.");
         System.out.println(conversation2.ask("What's your job?").content());
-        // Output: As a software engineer, my job involves designing, developing, testing, and maintaining software systems and applications. It can involve tasks such as writing code, debugging programs, troubleshooting issues, and collaborating with other team members to ensure the overall functionality and efficiency of the software being developed. I may also need to work on improving existing software, conducting research to identify new technologies or methods that could benefit my team, and keeping up with industry trends and best practices to continuously improve my skills and knowledge.
         System.out.println(conversation2.ask("What's your day to day work?").content());
-        // Output: As an AI language model, I do not have a physical day-to-day work environment. However, as a software engineer, a typical day may involve various activities such as:
+
+        // Load a conversation by the ID
+        final var conversation3 = chatGPT.loadConversation(conversation.id());
+        conversation3.ask("What did I ask you?");
+        // Should print the same as the first conversation
     }
 }
 ```

chatgpt-cli/src/main/java/io/github/kezhenxu94/chatgpt/cli/commands/ConversationCommands.java

@@ -7,14 +7,26 @@
 
 @ShellComponent
 public class ConversationCommands {
-  private final Conversation conversation;
+  private final ChatGPT chatGPT;
+
+  private Conversation conversation;
 
   public ConversationCommands(ChatGPT chatGPT) {
-    conversation = chatGPT.newConversation();
+    this.chatGPT = chatGPT;
+  }
+
+  @ShellMethod("Load a conversation by ID")
+  public String load(String id) throws Exception {
+    conversation = chatGPT.loadConversation(id);
+    return "You are now at conversation: " + id;
   }
 
-  @ShellMethod(value = "Ask a question")
+  @ShellMethod("Ask a question")
   public String ask(String question) throws Exception {
+    if (conversation == null) {
+      conversation = chatGPT.newConversation();
+    }
+
     final var answer = conversation.ask(question);
     return answer.content().trim();
   }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/ChatGPT.java

@@ -1,14 +1,25 @@
 package io.github.kezhenxu94.chatgpt;
 
+import java.io.IOException;
+import java.nio.file.Path;
+
 public interface ChatGPT {
+  public static final String API_URL = "https://api.openai.com/v1/chat/completions";
+
   Conversation newConversation();
 
   Conversation newConversation(String system);
 
+  Conversation newConversationWithID(String id);
+
+  Conversation loadConversation(String id) throws IOException;
+
   String apiKey();
 
   int conversationSize();
 
+  Path dataPath();
+
   static Builder builder() {
     return new ChatGPTBuilder();
   }
@@ -18,6 +29,10 @@ interface Builder {
 
     Builder apiKey(String apiKey);
 
+    Builder dataPath(String persistentPath);
+
+    Builder dataPath(Path persistentPath);
+
     ChatGPT build();
   }
 }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/ChatGPTBuilder.java

@@ -1,8 +1,12 @@
 package io.github.kezhenxu94.chatgpt;
 
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
 class ChatGPTBuilder implements ChatGPT.Builder {
   private String apiKey;
   private int conversationSize = Integer.MAX_VALUE;
+  private Path dataPath = Paths.get(System.getProperty("user.home"), ".chatgpt", "data");
 
   @Override
   public ChatGPT build() {
@@ -14,7 +18,7 @@ public ChatGPT build() {
       throw new IllegalArgumentException("API key is required");
     }
 
-    return new ChatGPTImpl(apiKey, conversationSize);
+    return new ChatGPTImpl(apiKey, conversationSize, dataPath);
   }
 
   public ChatGPTBuilder conversationSize(int conversationSize) {
@@ -27,4 +31,16 @@ public ChatGPT.Builder apiKey(String apiKey) {
     this.apiKey = apiKey;
     return this;
   }
+
+  @Override
+  public ChatGPT.Builder dataPath(String dataPath) {
+    this.dataPath = Paths.get(dataPath);
+    return this;
+  }
+
+  @Override
+  public ChatGPT.Builder dataPath(Path dataPath) {
+    this.dataPath = dataPath;
+    return this;
+  }
 }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/ChatGPTImpl.java

@@ -1,22 +1,65 @@
 package io.github.kezhenxu94.chatgpt;
 
+import io.github.kezhenxu94.chatgpt.message.Message;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+
 final class ChatGPTImpl implements ChatGPT {
+  private final List<Conversation> conversations;
+
   private final String apiKey;
   private final int conversationSize;
+  private final Path dataPath;
 
-  public ChatGPTImpl(String apiKey, int conversationSize) {
+  public ChatGPTImpl(String apiKey, int conversationSize, Path dataPath) {
     this.apiKey = apiKey;
     this.conversationSize = conversationSize;
+    this.dataPath = dataPath;
+
+    this.conversations = new ArrayList<>();
+
+    Runtime.getRuntime()
+        .addShutdownHook(
+            new Thread(
+                () -> {
+                  try {
+                    save();
+                  } catch (IOException e) {
+                    throw new RuntimeException(e);
+                  }
+                }));
   }
 
   @Override
   public Conversation newConversation() {
-    return new ConversationImpl(this);
+    final var newConversation = new ConversationImpl(this);
+    conversations.add(newConversation);
+    return newConversation;
   }
 
   @Override
   public Conversation newConversation(String system) {
-    return new ConversationImpl(this, system);
+    final var newConversation = new ConversationImpl(this, Message.ofSystem(system));
+    conversations.add(newConversation);
+    return newConversation;
+  }
+
+  @Override
+  public Conversation newConversationWithID(String id) {
+    final var conversation = new ConversationImpl(this, id);
+    conversations.add(conversation);
+    return conversation;
+  }
+
+  @Override
+  public Conversation loadConversation(String id) throws IOException {
+    final var conversation = new ConversationImpl(this, id);
+    conversations.add(conversation);
+
+    conversation.load();
+    return conversation;
   }
 
   @Override
@@ -28,4 +71,15 @@ public String apiKey() {
   public int conversationSize() {
     return conversationSize;
   }
+
+  @Override
+  public Path dataPath() {
+    return dataPath;
+  }
+
+  void save() throws IOException {
+    for (final var conversation : conversations) {
+      conversation.save();
+    }
+  }
 }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/Conversation.java

@@ -2,8 +2,16 @@
 
 import io.github.kezhenxu94.chatgpt.message.Message;
 
-import java.io.IOException;
+import java.io.IOException;import java.util.List;
 
 public interface Conversation {
   Message ask(String question) throws IOException, InterruptedException;
+
+  String id();
+
+  List<Message> messages();
+
+  void save() throws IOException;
+
+  void load() throws IOException;
 }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/ConversationImpl.java

@@ -1,37 +1,43 @@
 package io.github.kezhenxu94.chatgpt;
 
+import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.github.kezhenxu94.chatgpt.internal.ChatGPTHttpRequest;
 import io.github.kezhenxu94.chatgpt.internal.ChatGPTHttpResponse;
 import io.github.kezhenxu94.chatgpt.internal.JsonBodyHandler;
 import io.github.kezhenxu94.chatgpt.message.Message;
-
+import io.github.kezhenxu94.chatgpt.message.SystemMessage;
 import java.io.IOException;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.time.Duration;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.UUID;
 
 final class ConversationImpl implements Conversation {
-  private static final String API_URL = "https://api.openai.com/v1/chat/completions";
-
   private final ChatGPT chatGPT;
+  private final String uuid;
   private final List<Message> messages;
   private final HttpClient httpClient;
   private final ObjectMapper objectMapper;
 
-  ConversationImpl(ChatGPT chatGPT) {
+  ConversationImpl(ChatGPT chatGPT, String uuid) {
     this.chatGPT = chatGPT;
+    this.uuid = uuid;
     this.messages = new LinkedList<>();
-    this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build();
+    this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build();
     this.objectMapper = new ObjectMapper();
   }
 
-  ConversationImpl(ChatGPT chatGPT, String system) {
+  ConversationImpl(ChatGPT chatGPT) {
+    this(chatGPT, UUID.randomUUID().toString());
+  }
+
+  ConversationImpl(ChatGPT chatGPT, SystemMessage system) {
     this(chatGPT);
-    messages.add(Message.ofSystem(system));
+    messages.add(system);
   }
 
   @Override
@@ -47,7 +53,7 @@ public Message ask(String question) throws IOException, InterruptedException {
     final var response =
         httpClient.send(
             HttpRequest.newBuilder()
-                .uri(URI.create(API_URL))
+                .uri(URI.create(ChatGPT.API_URL))
                 .header("Authorization", "Bearer " + chatGPT.apiKey())
                 .header("Content-Type", "application/json")
                 .POST(HttpRequest.BodyPublishers.ofString(objectMapper.writeValueAsString(request)))
@@ -60,4 +66,42 @@ public Message ask(String question) throws IOException, InterruptedException {
 
     return choice.message();
   }
+
+  @Override
+  public String id() {
+    return uuid;
+  }
+
+  @Override
+  public List<Message> messages() {
+    return messages;
+  }
+
+  @Override
+  public void save() throws IOException {
+    final var dataPath = chatGPT.dataPath();
+    if (!dataPath.toFile().exists()) {
+      //noinspection ResultOfMethodCallIgnored
+      dataPath.toFile().mkdirs();
+    }
+
+    final var conversationFile = dataPath.resolve(id() + ".json").toFile();
+    objectMapper.writeValue(conversationFile, messages());
+  }
+
+  @Override
+  public void load() throws IOException {
+    final var dataPath = chatGPT.dataPath();
+    if (!dataPath.toFile().exists()) {
+      return;
+    }
+    final var conversationFile = dataPath.resolve(id() + ".json").toFile();
+    if (!conversationFile.exists()) {
+      return;
+    }
+
+    final var messages =
+        objectMapper.readValue(conversationFile, new TypeReference<List<Message>>() {});
+    this.messages.addAll(messages);
+  }
 }

chatgpt-java-sdk/src/main/java/io/github/kezhenxu94/chatgpt/message/Message.java

@@ -17,11 +17,11 @@ public interface Message {
   @JsonProperty
   String content();
 
-  static Message ofUser(String content) {
+  static UserMessage ofUser(String content) {
     return new UserMessage(content);
   }
 
-  static Message ofSystem(String content) {
+  static SystemMessage ofSystem(String content) {
     return new SystemMessage(content);
   }
 }

chatgpt-java-sdk/src/test/java/io/github/kezhenxu94/chatgpt/ChatGPTTest.java

@@ -1,18 +1,29 @@
 package io.github.kezhenxu94.chatgpt;
 
-import java.io.IOException;
+import java.io.IOException;import java.nio.file.Files;
 
 public class ChatGPTTest {
   public static void main(String[] args) throws IOException, InterruptedException {
-    final var chatGPT = ChatGPT.builder().conversationSize(2).build();
+    final var chatGPT = ChatGPT
+            .builder()
+            .dataPath(Files.createTempDirectory("chatgpt")) // Persist the chat history to a data path
+            .build();
+
+    // Start a new conversation
     final var conversation = chatGPT.newConversation();
     System.out.println(conversation.ask("What's your name?").content());
-    System.out.println(conversation.ask("What's your hobbit?").content());
-    System.out.println(conversation.ask("What's your favorite song?").content());
+    // Output: I'm an AI language model developed by OpenAI, and I don't have a name. What can I help you with today?
     System.out.println(conversation.ask("What did I ask you?").content());
+    // Output: You asked for my name.
+    conversation.save(); // Save the history manually, conversations are saved on shutdown by default.
 
     final var conversation2 = chatGPT.newConversation("You are a software engineer.");
     System.out.println(conversation2.ask("What's your job?").content());
     System.out.println(conversation2.ask("What's your day to day work?").content());
+
+    // Load a conversation by the ID
+    final var conversation3 = chatGPT.loadConversation(conversation.id());
+    conversation3.ask("What did I ask you?");
+    // Should print the same as the first conversation
   }
 }