From ef2bd7285dc44446e8adec485d4ea5532387c9c8 Mon Sep 17 00:00:00 2001
From: exactchange <benny@exactchange.network>
Date: Wed, 20 Mar 2024 11:54:17 -0700
Subject: [PATCH] make chat function stateful

---
 arthas-api/api/post/prompt/index.js | 35 ++++++++++++++++++++---------
 arthas-api/index.js                 |  3 ++-
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/arthas-api/api/post/prompt/index.js b/arthas-api/api/post/prompt/index.js
index 234e59a..673952c 100644
--- a/arthas-api/api/post/prompt/index.js
+++ b/arthas-api/api/post/prompt/index.js
@@ -28,11 +28,15 @@ const {
 
 const { prefixInput } = require('arthasgpt/src/utils/prefix');
 
+const __statefulChatFunction = {};
+
 /**
  * Prompt
  */
 
 module.exports = asyncCache => async (req, res) => {
+  const config = await asyncCache.getItem('config');
+
   const timeout = await asyncCache.getItem('timeout');
 
   let answer = await asyncCache.getItem('answer');
@@ -44,8 +48,6 @@ module.exports = asyncCache => async (req, res) => {
       body.push(chunk);
     })
     .on('end', async () => {
-      const config = await asyncCache.getItem('config');
-
       body = Buffer.concat(body).toString();
 
       const { key, input } = JSON.parse(body || '{}');
@@ -112,16 +114,29 @@ module.exports = asyncCache => async (req, res) => {
         await delay(timeout);
       }
 
-      if (isVerbose) {
-        log(CREATING_AGENT);
-      }
+      const currentChat = __statefulChatFunction?.[key];
 
-      answer = await ArthasGPT({
-        ...currentConfig,
+      if (currentChat) {
+        await currentChat(messageResponse);
+      } else {
+        if (isVerbose) {
+          log(CREATING_AGENT);
+        }
 
-        query: messageResponse,
-        cache: false
-      });
+        answer = await ArthasGPT({
+          ...currentConfig,
+
+          query: messageResponse,
+          cache: true
+        });
+
+        await asyncCache.setItem('answer', answer);
+
+        // Cache the stateful chat method in the API
+        // (antipattern)
+
+        __statefulChatFunction[key] = answer?.chat;
+      }
 
       res.end(JSON.stringify({
         success: true,
diff --git a/arthas-api/index.js b/arthas-api/index.js
index f45b134..e0e353a 100644
--- a/arthas-api/index.js
+++ b/arthas-api/index.js
@@ -30,7 +30,8 @@ const PORT = 8000;
  * Config
  */
 
-const numCPUs = availableParallelism();
+// const numCPUs = availableParallelism();
+const numCPUs = 1;
 
 const onCluster = require('./events/cluster');
 const onWorker = require('./events/worker');