[lib] Add conversation serialization (#13)

maxijonson · May 1, 2023 · 7cea741 · 7cea741
1 parent 7018d78
commit 7cea741
Show file tree

Hide file tree

Showing 25 changed files with 410 additions and 489 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/cli/src/contexts/providers/PersistenceProvider.tsx b/packages/cli/src/contexts/providers/PersistenceProvider.tsx
@@ -1,7 +1,5 @@
 import React from "react";
 import fs from "fs";
-import { Persistence, persistenceSchema } from "../../entities/persistence.js";
-import { PersistenceMessage } from "../../entities/persistenceMessage.js";
 import useConversationManager from "../../hooks/useConversationManager.js";
 import {
     PersistenceContext,
@@ -20,19 +18,7 @@ export default ({ children }: PersistenceProviderProps) => {
 
     const save = React.useCallback(() => {
         if (!conversation || !saveFile) return;
-        const persistedConversations: Persistence = {
-            conversation: {
-                ...conversation.getConfig(),
-                messages: conversation.getMessages().map(
-                    (message): PersistenceMessage => ({
-                        content: message.content,
-                        role: message.role,
-                    })
-                ),
-            },
-        };
-        const parsed = persistenceSchema.parse(persistedConversations);
-        const json = JSON.stringify(parsed);
+        const json = JSON.stringify(conversation.toJSON());
         fs.writeFileSync(saveFile, json);
     }, [conversation, saveFile]);
 
@@ -58,45 +44,7 @@ export default ({ children }: PersistenceProviderProps) => {
             try {
                 const data = fs.readFileSync(loadFile, "utf8");
                 const json = JSON.parse(data);
-                const { conversation: persistedConversation } =
-                    persistenceSchema.parse(json);
-                const { messages, disableModeration, ...config } =
-                    persistedConversation;
-                const newConversation = new Conversation({
-                    ...config,
-                    disableModeration: true,
-                });
-
-                for (const message of messages) {
-                    try {
-                        switch (message.role) {
-                            case "user":
-                                await newConversation.addUserMessage(
-                                    message.content
-                                );
-                                break;
-                            case "assistant":
-                                await newConversation.addAssistantMessage(
-                                    message.content
-                                );
-                                break;
-                            case "system":
-                                newConversation.setContext(message.content);
-                        }
-                    } catch (e) {
-                        console.error(
-                            "Error while loading message",
-                            (e as Error).message
-                        );
-                    }
-                }
-
-                newConversation.setConfig(
-                    {
-                        disableModeration,
-                    },
-                    true
-                );
+                const newConversation = await Conversation.fromJSON(json);
                 setConversation(newConversation);
             } catch (e) {
                 console.error(e);
@@ -117,10 +65,7 @@ export default ({ children }: PersistenceProviderProps) => {
             messageOffs.push(
                 message.onMessageStreamingStop(() => {
                     save();
-                })
-            );
-
-            messageOffs.push(
+                }),
                 message.onMessageUpdate((_, m) => {
                     if (m.isStreaming) return;
                     save();

diff --git a/packages/cli/src/entities/persistence.ts b/packages/cli/src/entities/persistence.ts
diff --git a/packages/cli/src/entities/persistenceConversation.ts b/packages/cli/src/entities/persistenceConversation.ts
diff --git a/packages/cli/src/entities/persistenceMessage.ts b/packages/cli/src/entities/persistenceMessage.ts
diff --git a/packages/cli/src/index.tsx b/packages/cli/src/index.tsx
@@ -132,7 +132,9 @@ const saveFile = (() => {
     const savePath = path.normalize(
         typeof save === "string"
             ? save
-            : `gptturbo-conversation-${new Date().toISOString()}`
+            : `gptturbo-conversation-${new Date()
+                  .toISOString()
+                  .replace(/:/g, "-")}`
     );
     return savePath.endsWith(".json") ? savePath : `${savePath}.json`;
 })();

diff --git a/packages/discord/src/managers/ConversationManager.ts b/packages/discord/src/managers/ConversationManager.ts
@@ -1,4 +1,4 @@
-import { ChatCompletionRequestMessageRoleEnum, Conversation } from "gpt-turbo";
+import { Conversation, MessageModel } from "gpt-turbo";
 import getConversationConfig from "../utils/getConversationConfig.js";
 import BotException from "../exceptions/BotException.js";
 import QuotaManager from "./QuotaManager.js";
@@ -8,10 +8,7 @@ export default class ConversationManager {
     constructor(private readonly quotaManager: QuotaManager) {}
 
     public async getChatCompletion(
-        messages: (
-            | { content: string; role: ChatCompletionRequestMessageRoleEnum }
-            | string
-        )[],
+        messages: (MessageModel | string)[],
         userId: string
     ) {
         const user: ConversationUser = `discord-${userId}`;
@@ -52,15 +49,9 @@ export default class ConversationManager {
     }
 
     private getAlternatedMessages(
-        messages: (
-            | { content: string; role: ChatCompletionRequestMessageRoleEnum }
-            | string
-        )[]
-    ): { content: string; role: ChatCompletionRequestMessageRoleEnum }[] {
-        const alternatedMessages: {
-            content: string;
-            role: ChatCompletionRequestMessageRoleEnum;
-        }[] = [];
+        messages: (MessageModel | string)[]
+    ): MessageModel[] {
+        const alternatedMessages: MessageModel[] = [];
 
         messages.forEach((message, index) => {
             alternatedMessages.push({

diff --git a/packages/discord/src/message-handlers/DMReplyHandler.ts b/packages/discord/src/message-handlers/DMReplyHandler.ts
@@ -3,7 +3,7 @@ import MessageHandler from "./MessageHandler.js";
 import getPromptAndReplyMessages from "../utils/getPromptAndReplyMessages.js";
 import BotException from "../exceptions/BotException.js";
 import getCleanContent from "../utils/getCleanContent.js";
-import { ChatCompletionRequestMessageRoleEnum } from "gpt-turbo";
+import { MessageModel } from "gpt-turbo";
 
 export default class DMReplyHandler extends MessageHandler {
     public get name(): string {
@@ -41,12 +41,7 @@ export default class DMReplyHandler extends MessageHandler {
         const [conversationMessages, prompt] = await Promise.all([
             Promise.all(
                 messages.map(
-                    async (
-                        m
-                    ): Promise<{
-                        content: string;
-                        role: ChatCompletionRequestMessageRoleEnum;
-                    }> => ({
+                    async (m): Promise<MessageModel> => ({
                         content: (await getCleanContent(m)) || "Hello",
                         role:
                             m.author.id === message.author.id

diff --git a/packages/lib/package.json b/packages/lib/package.json
@@ -71,7 +71,8 @@
     },
     "dependencies": {
         "gpt-token-utils": "^1.2.0",
-        "uuid": "^9.0.0"
+        "uuid": "^9.0.0",
+        "zod": "^3.21.4"
     },
     "gitHead": "328d04dc73d86236dba158a181581a53caf37720"
 }