bugfix on chatting API

Muimi-Chat · Jun 24, 2024 · d1974dc · d1974dc
1 parent 5eaab53
commit d1974dc
Show file tree

Hide file tree

Showing 14 changed files with 125 additions and 68 deletions.
diff --git a/.env.example b/.env.example
@@ -25,6 +25,9 @@ DEFAULT_STARTING_TOKEN=25000
 # Database #
 ############
 
+# For database explorer to connect to.
+POSTGRESQL_PORT=44902
+
 # This is for postgreSQL connection credentials
 POSTGRES_USERNAME=test
 POSTGRES_PASSWORD=test
@@ -41,5 +44,5 @@ QDRANT_API_KEY=asdasdasd
 
 # For communication with User service; Should be same as whats set in user service's
 USER_SERVICE_API_TOKEN=abcdefg
-USER_SERVICE_API_DOMAIN=localhost:44801
+USER_SERVICE_API_DOMAIN=user-api:8000
 USER_SERVICE_SSL_ENABLED=false
diff --git a/docker-compose-prod.yml b/docker-compose-prod.yml
@@ -30,6 +30,8 @@ services:
       - chat-db
       - chat-vector-db
       - chat-cache
+    networks:
+      - muimi_container_network  
 
   # database
   chat-db:
@@ -39,19 +41,31 @@ services:
       POSTGRES_USER: ${POSTGRES_USERNAME}
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
       POSTGRES_DB: ${POSTGRES_MAIN_DB}
+    ports:
+      - "${POSTGRESQL_PORT}:5432"
     volumes:
       - ./pgdata:/var/lib/postgresql/data:z
+    networks:
+      - muimi_container_network  
 
   chat-vector-db:
     image: qdrant/qdrant:v1.9.4
     environment:
       QDRANT__SERVICE__API_KEY: ${QDRANT_API_KEY}
     volumes:
       - "./qdrant_storage:/qdrant/storage:z"
+    networks:
+      - muimi_container_network  
 
   chat-cache:
     image: redis:7.2.4
     command: ["redis-server", "--requirepass", "${REDIS_PASSWORD}"]
     # Runs on 6379 Docker
     volumes:
       - ./redisdata:/data:z
+    networks:
+      - muimi_container_network  
+
+networks:
+  muimi_container_network:
+    external: true
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -28,13 +28,17 @@ services:
       - chat-db
       - chat-vector-db
       - chat-cache
+    networks:
+      - muimi_container_network  
 
   chat-vector-db:
     image: qdrant/qdrant:v1.9.4
     environment:
       QDRANT__SERVICE__API_KEY: ${QDRANT_API_KEY}
     volumes:
       - "./qdrant_storage:/qdrant/storage:z"
+    networks:
+      - muimi_container_network  
 
   # database
   chat-db:
@@ -44,13 +48,22 @@ services:
       POSTGRES_USER: ${POSTGRES_USERNAME}
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
       POSTGRES_DB: ${POSTGRES_MAIN_DB}
+    ports:
+      - "${POSTGRESQL_PORT}:5432"
     volumes:
       - ./pgdata:/var/lib/postgresql/data:z
+    networks:
+      - muimi_container_network  
 
   chat-cache:
     image: redis:7.2.4
     command: ["redis-server", "--requirepass", "${REDIS_PASSWORD}"]
     # Runs on 6379 Docker
     volumes:
       - ./redisdata:/data:z
+    networks:
+      - muimi_container_network  
 
+networks:
+  muimi_container_network:
+    external: true
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -18,13 +18,15 @@
   "dependencies": {
     "@qdrant/js-client-rest": "^1.9.0",
     "@types/crypto-js": "^4.2.2",
+    "@types/express-ws": "^3.0.4",
     "@types/ws": "^8.5.10",
     "argon2": "^0.40.3",
     "axios": "^1.7.2",
     "cors": "^2.8.5",
     "crypto-js": "^4.2.0",
     "drizzle-orm": "^0.30.10",
     "express": "^4.19.2",
+    "express-ws": "^5.0.2",
     "openai": "^4.47.3",
     "pg": "^8.11.5",
     "postgres": "^3.4.4",

diff --git a/src/api/consumers/chatMessageConsumer.ts b/src/api/consumers/chatMessageConsumer.ts
@@ -1,4 +1,4 @@
-import { WebSocket } from 'ws';
+import { RawData, WebSocket } from 'ws';
 import insertLog from '../repositories/insertLog';
 
 /**
@@ -36,9 +36,9 @@ function getSenderUserUUID(data: any): Promise<string | null> {
 }
 
 // TODO: Support custom system messages...
-export default async function chatMessageConsumer(socketClient: WebSocket, content: string) {
+export default async function chatMessageConsumer(socketClient: WebSocket, content: RawData) {
     try {
-        const data = JSON.parse(content)
+        const data = JSON.parse(content.toString())
 
         const userUUID = await getSenderUserUUID(data)
         if (userUUID == null) {
@@ -146,26 +146,27 @@ export default async function chatMessageConsumer(socketClient: WebSocket, conte
         // Start streaming in the message chunk by chunk
         const botReplyStream = await chatStream(openAiClient, messageChain, chatModel, userUUID)
         for await (const chunk of botReplyStream) {
-            const chunkContent = chunk.choices[0].delta.content
-            botReply += chunkContent
-            socketClient.send(JSON.stringify({ 
-                status: "CHUNK",
-                message: "This is a chunk of the bot's message...",
-                chunk_content: chunkContent
-            }));
-
-            // Last chunk, fetch the token cost.
-            if (chunk.choices[0].finish_reason == "stop") {
+            if (chunk.usage != null) {
                 promptTokenCost = chunk.usage!.prompt_tokens
                 completionTokenCost = chunk.usage!.completion_tokens
+                // Because a chunk with usage has no message...
+                continue; 
             }
-        }
 
-        // TODO: Sequencial undo in case of error....
+            if (chunk.choices[0].finish_reason != "stop") {
+                const chunkContent = chunk.choices[0].delta.content
+                botReply += chunkContent
+                socketClient.send(JSON.stringify({ 
+                    status: "CHUNK",
+                    message: "This is a chunk of the bot's message...",
+                    chunk_content: chunkContent
+                }));
+            }
+        }
 
         let userMessageObject;
         let botMessageObject
-        await db.transaction(async (tsx) => {
+        await db.transaction(async () => {
             if (!userAccountModel.freeTokenUsage) {
                 await updateAccountTokenByUUID(userUUID, userAccountModel.token - promptTokenCost - completionTokenCost)
             }

diff --git a/src/api/controllers/getConversationMessages.ts b/src/api/controllers/getConversationMessages.ts
@@ -51,7 +51,7 @@ export default async function getConversationMessages(req: Request, res: Respons
             "WARNING"
           );
           return res.status(401).json({
-            status: "BAD_TOKEN",
+            status: "BAD_SESSION",
             message: "Session token is bad, relogin!",
           });
         }

diff --git a/src/api/controllers/getTokenCountController.ts b/src/api/controllers/getTokenCountController.ts
@@ -12,6 +12,7 @@ export default async function getTokenCountController(req: Request, res: Respons
         url: REDIS_CONNECTION_STRING
     });
 
+    console.log("Connecting to Redis");
     client.on('error', async (err: any) => {
         console.log('Connecting Redis Client Error', err);
         await insertLog(`Error connecting to Redis :: ${err}`);
@@ -32,20 +33,22 @@ export default async function getTokenCountController(req: Request, res: Respons
                     message: 'Missing required parameters'
                 });
             }
-
             // See if token exists in cache, if not, re-validate from user service,
             // then flag token as valid for 1 minute.
             let userUUID = await client.get(`${username}_${sessionToken}`);
             if (userUUID == null) {
+                console.log("User UUID not found in cache, revalidating from API");
                 let userInformation = await fetchUserInformation(sessionToken, userAgent, username)
                 if (userInformation.status != "SUCCESS") {
+                    console.log(`Failed to fetch user information with status :: ${userInformation.status}`)
                     await insertLog(`Failed to fetch user information with status :: ${userInformation.status}`, "WARNING")
                     return res.status(401).json({
-                        status: 'BAD_TOKEN',
+                        status: 'BAD_SESSION',
                         message: 'Session token is bad, relogin!'
                     });
                 }
 
+                console.log(`User UUID Revalidated :: ${userInformation.uuid}`);
                 userUUID = userInformation.uuid;
                 await client.set(`${username}_${sessionToken}`, userUUID, {
                     EX: 60,

diff --git a/src/api/helpers/trimString.ts b/src/api/helpers/trimString.ts
@@ -1,3 +1,3 @@
-function trimString(str: string, maxLength: number): string {
+export default function trimString(str: string, maxLength: number): string {
     return str.length > maxLength ? str.substring(0, maxLength) : str;
 }
diff --git a/src/api/repositories/selectConversationsByUserID.ts b/src/api/repositories/selectConversationsByUserID.ts
@@ -10,7 +10,7 @@ export default async function selectConversationsByUserID(userID: string, select
     .where(and(
             eq(conversation.accountID, userID),
             eq(conversation.deleted, selectDeleted),
-            eq(conversation.archived, selectArchived)
+            selectArchived ? undefined : eq(conversation.archived, false)
         )  
     )
     .orderBy(desc(conversation.creationDate));