From fd9d6fcb0356b2a93df4874e910afce6a7482f0b Mon Sep 17 00:00:00 2001
From: Jeffrey-Wang <admin@wj2015.com>
Date: Sun, 11 Jun 2023 23:30:26 +0800
Subject: [PATCH] fix: zep memory

---
 .../nodes/memory/ZepMemory/ZepMemory.ts       | 14 +++++++++--
 packages/server/src/ChildProcess.ts           |  4 +++-
 packages/server/src/index.ts                  | 23 +++++++++++--------
 .../ui/src/views/chatmessage/ChatMessage.js   |  2 +-
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/packages/components/nodes/memory/ZepMemory/ZepMemory.ts b/packages/components/nodes/memory/ZepMemory/ZepMemory.ts
index 95b83bbb..cf1d8e58 100644
--- a/packages/components/nodes/memory/ZepMemory/ZepMemory.ts
+++ b/packages/components/nodes/memory/ZepMemory/ZepMemory.ts
@@ -35,6 +35,14 @@ class ZepMemory_Memory implements INode {
                 type: 'boolean',
                 default: true
             },
+            {
+                label: 'Session Id',
+                name: 'sessionId',
+                type: 'string',
+                placeholder: 'if empty, chatId will be used automatically',
+                default: '',
+                additionalParams: true
+            },
             {
                 label: 'Auto Summary Template',
                 name: 'autoSummaryTemplate',
@@ -88,11 +96,13 @@ class ZepMemory_Memory implements INode {
         const inputKey = nodeData.inputs?.inputKey as string
         const autoSummaryTemplate = nodeData.inputs?.autoSummaryTemplate as string
         const autoSummary = nodeData.inputs?.autoSummary as boolean
-        const sessionId = options?.chatId as string
+        const sessionId = nodeData.inputs?.sessionId as string
+
+        const chatId = options?.chatId as string
 
         const obj: ZepMemoryInput = {
             baseURL,
-            sessionId,
+            sessionId: sessionId ? sessionId : chatId,
             aiPrefix,
             humanPrefix,
             returnMessages: true,
diff --git a/packages/server/src/ChildProcess.ts b/packages/server/src/ChildProcess.ts
index e5f95170..07b52909 100644
--- a/packages/server/src/ChildProcess.ts
+++ b/packages/server/src/ChildProcess.ts
@@ -1,5 +1,6 @@
 import { IChildProcessMessage, IReactFlowNode, IReactFlowObject, IRunChatflowMessageValue, INodeData } from './Interface'
 import { buildLangchain, constructGraphs, getEndingNode, getStartingNodes, resolveVariables } from './utils'
+import { getChatId } from './index'
 
 export class ChildProcess {
     /**
@@ -76,6 +77,7 @@ export class ChildProcess {
             const { startingNodeIds, depthQueue } = getStartingNodes(nonDirectedGraph, endingNodeId)
 
             /*** BFS to traverse from Starting Nodes to Ending Node ***/
+            const chatId = await getChatId(chatflow.id)
             const reactFlowNodes = await buildLangchain(
                 startingNodeIds,
                 nodes,
@@ -83,7 +85,7 @@ export class ChildProcess {
                 depthQueue,
                 componentNodes,
                 incomingInput.question,
-                '',
+                chatId,
                 incomingInput?.overrideConfig
             )
 
diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts
index 5e91a30a..78be231a 100644
--- a/packages/server/src/index.ts
+++ b/packages/server/src/index.ts
@@ -506,15 +506,8 @@ export class App {
             })
             if (!chatflow) return res.status(404).send(`Chatflow ${chatflowid} not found`)
 
-            // first chatmessage id as the unique chat id
-            const firstChatMessage = await this.AppDataSource.getRepository(ChatMessage)
-                .createQueryBuilder('cm')
-                .select('cm.id')
-                .where('chatflowid = :chatflowid', { chatflowid })
-                .orderBy('cm.createdDate', 'ASC')
-                .getOne()
-            if (!firstChatMessage) return res.status(500).send(`Chatflow ${chatflowid} first message not found`)
-            const chatId = firstChatMessage.id
+            const chatId = await getChatId(chatflow.id)
+            if (!chatId) return res.status(500).send(`Chatflow ${chatflowid} first message not found`)
 
             if (!isInternal) {
                 await this.validateKey(req, res, chatflow)
@@ -672,6 +665,18 @@ export class App {
     }
 }
 
+export async function getChatId(chatflowid: string) {
+    // first chatmessage id as the unique chat id
+    const firstChatMessage = await getDataSource()
+        .getRepository(ChatMessage)
+        .createQueryBuilder('cm')
+        .select('cm.id')
+        .where('chatflowid = :chatflowid', { chatflowid })
+        .orderBy('cm.createdDate', 'ASC')
+        .getOne()
+    return firstChatMessage ? firstChatMessage.id : ''
+}
+
 let serverApp: App | undefined
 
 export async function start(): Promise<void> {
diff --git a/packages/ui/src/views/chatmessage/ChatMessage.js b/packages/ui/src/views/chatmessage/ChatMessage.js
index 2f5415d0..5021cd9b 100644
--- a/packages/ui/src/views/chatmessage/ChatMessage.js
+++ b/packages/ui/src/views/chatmessage/ChatMessage.js
@@ -118,7 +118,7 @@ export const ChatMessage = ({ open, chatflowid, isDialog }) => {
 
         setLoading(true)
         setMessages((prevMessages) => [...prevMessages, { message: userInput, type: 'userMessage' }])
-        // waiting for first chatmessage uploaded, the first chatmessage id will be chatId for every components
+        // waiting for first chatmessage saved, the first chatmessage will be used in sendMessageAndGetPrediction
         await addChatMessage(userInput, 'userMessage')
 
         // Send user question and history to API