Added download chat history button for AI Assistant, including chat m…

…essage DB object, db functions, and http routing for downloading csv
encorelab · Oct 28, 2024 · ba03592 · ba03592
1 parent e0d1a86
commit ba03592
Show file tree

Hide file tree

Showing 8 changed files with 184 additions and 30 deletions.
diff --git a/backend/src/api/chatHistory.ts b/backend/src/api/chatHistory.ts
@@ -0,0 +1,45 @@
+// backend/src/api/chatHistory.ts
+
+import { Router } from 'express';
+import dalChatMessage from '../repository/dalChatMessage';
+
+const router = Router();
+
+router.post('/', async (req, res) => {
+  try {
+    const { boardId, userId } = req.body;
+    const chatHistory = await dalChatMessage.getByBoardIdAndUserId(
+      boardId,
+      userId
+    );
+
+    // Format the chat history as a CSV string
+    const csvString = formatChatHistoryAsCSV(chatHistory);
+
+    res.setHeader('Content-Type', 'text/csv');
+    res.setHeader(
+      'Content-Disposition',
+      'attachment; filename="chat_history.csv"'
+    );
+    res.send(csvString);
+  } catch (error) {
+    console.error('Error fetching or formatting chat history:', error);
+    res.status(500).json({ error: 'Failed to download chat history.' });
+  }
+});
+
+function formatChatHistoryAsCSV(chatHistory: any[]): string {
+  // 1. Create header row
+  let csvString = 'Timestamp,Role,Content\n';
+
+  // 2. Add data rows
+  for (const message of chatHistory) {
+    csvString += `"${message.createdAt.toLocaleString()}","${
+      message.role
+    }","${message.content.replace(/"/g, '""')}"\n`;
+  }
+
+  return csvString;
+}
+
+export default router;
diff --git a/backend/src/models/ChatMessage.ts b/backend/src/models/ChatMessage.ts
@@ -0,0 +1,21 @@
+import { prop, getModelForClass, modelOptions } from '@typegoose/typegoose';
+
+@modelOptions({
+  schemaOptions: { collection: 'chatMessages', timestamps: true },
+})
+export class ChatMessageModel {
+  @prop({ required: true })
+  public userId!: string;
+
+  @prop({ required: true })
+  public boardId!: string;
+
+  @prop({ required: true, enum: ['user', 'assistant'] })
+  public role!: string;
+
+  @prop({ required: true })
+  public content!: string;
+}
+
+const ChatMessage = getModelForClass(ChatMessageModel);
+export default ChatMessage;
diff --git a/backend/src/repository/dalChatMessage.ts b/backend/src/repository/dalChatMessage.ts
@@ -0,0 +1,31 @@
+import ChatMessage, { ChatMessageModel } from '../models/ChatMessage';
+
+export const save = async (message: ChatMessageModel) => {
+  try {
+    const savedMessage = await ChatMessage.create(message);
+    return savedMessage;
+  } catch (err) {
+    throw new Error(JSON.stringify(err, null, ' '));
+  }
+};
+
+export const getByBoardIdAndUserId = async (
+  boardId: string,
+  userId: string
+) => {
+  try {
+    const messages = await ChatMessage.find({ boardId, userId }).sort({
+      createdAt: 1,
+    }); // Sort by createdAt in descending order (oldest to most recent)
+    return messages;
+  } catch (err) {
+    throw new Error(JSON.stringify(err, null, ' '));
+  }
+};
+
+const dalChatMessage = {
+  save,
+  getByBoardIdAndUserId,
+};
+
+export default dalChatMessage;
diff --git a/backend/src/server.ts b/backend/src/server.ts
@@ -22,6 +22,7 @@ import learner from './api/learner';
 import { isAuthenticated } from './utils/auth';
 import RedisClient from './utils/redis';
 import aiRouter from './api/ai';
+import chatHistoryRouter from './api/chatHistory'; 
 dotenv.config();
 
 const port = process.env.PORT || 8001;
@@ -68,6 +69,7 @@ app.use('/api/trace', isAuthenticated, trace);
 app.use('/api/todoItems', isAuthenticated, todoItems);
 app.use('/api/learner', isAuthenticated, learner);
 app.use('/api/ai', isAuthenticated, aiRouter);
+app.use('/api/chat-history', chatHistoryRouter);
 
 app.get('*', (req, res) => {
   res.sendFile(path.join(staticFilesPath, 'index.html'));

diff --git a/backend/src/services/vertexAI/index.ts b/backend/src/services/vertexAI/index.ts
@@ -1,5 +1,6 @@
 import dalVote from '../../repository/dalVote';
 import dalBucket from '../../repository/dalBucket';
+import dalChatMessage from '../../repository/dalChatMessage';
 import { BucketModel } from '../../models/Bucket';
 import { generateUniqueID } from '../../utils/Utils';
 import dalPost from '../../repository/dalPost';
@@ -288,6 +289,14 @@ function removeJsonMarkdown(text: string): string {
   return text.replace(pattern, '$1');
 }
 
+function removeEnd(str: string): string {
+  const endIndex = str.indexOf('<END>');
+  if (endIndex !== -1) {
+    return str.substring(0, endIndex);
+  }
+  return str; // Return original string if <END> is not found
+}
+
 async function sendMessage(
   posts: any[],
   prompt: string,
@@ -297,6 +306,17 @@ async function sendMessage(
     // Send an initial acknowledgment
     socket.emit(SocketEvent.AI_RESPONSE, { status: 'Received' });
 
+    const userId = socket.data.userId;
+    const boardId = socket.data.boardId;
+
+    // Save user prompt
+    await dalChatMessage.save({
+      userId: userId, 
+      boardId: boardId,
+      role: 'user',
+      content: prompt 
+    });
+
     // 1. Fetch Upvote Counts and Create Map
     const upvoteMap = await fetchUpvoteCounts(posts);
 
@@ -311,7 +331,7 @@ async function sendMessage(
     );
 
     // 4. Fetch and Format Buckets
-    const bucketsToSend = await fetchAndFormatBuckets(posts);
+    const bucketsToSend = await fetchAndFormatBuckets(boardId);
 
     // 5. Construct and Send Message to LLM (streaming)
     constructAndSendMessage(postsWithBucketIds, bucketsToSend, prompt).then(
@@ -337,7 +357,6 @@ async function sendMessage(
           // Async IIFE
           for await (const item of stream) {
             partialResponse += item.candidates[0].content.parts[0].text;
-            // console.log("Partial response:", partialResponse);
 
             socket.emit(SocketEvent.AI_RESPONSE, {
               status: 'Processing',
@@ -366,6 +385,15 @@ async function sendMessage(
               status: 'Completed',
               response: finalResponse.response,
             });
+
+            // Save AI response
+            await dalChatMessage.save({
+              userId: userId,
+              boardId: boardId,
+              role: 'assistant',
+              content: removeEnd(finalResponse.response)
+            }); 
+
           } else {
             const errorMessage =
               'Invalid response formatting. Please try again.\n\n' +
@@ -425,7 +453,7 @@ async function sendMessage(
 async function performDatabaseOperations(
   response: any,
   posts: any[],
-  socket: socketIO.Socket
+  socket: socketIO.Socket,
 ) {
   const validPostIds = new Set(posts.map((post) => post.postID));
 
@@ -435,7 +463,7 @@ async function performDatabaseOperations(
 
       const newBucket: BucketModel = {
         bucketID: generateUniqueID(),
-        boardID: posts[0].boardID,
+        boardID: socket.data.boardId,
         name: bucketName,
         posts: [],
       };
@@ -454,7 +482,7 @@ async function performDatabaseOperations(
 
       const newBucket: BucketModel = {
         bucketID: generateUniqueID(),
-        boardID: posts[0].boardID,
+        boardID: socket.data.boardId,
         name: name,
         posts: [],
       };
@@ -494,7 +522,7 @@ async function performDatabaseOperations(
             postID
           );
         } else if (name) {
-          const bucket = await dalBucket.getByName(name, posts[0].boardID);
+          const bucket = await dalBucket.getByName(name, socket.data.boardId);
           if (bucket) {
             actionsByBucket[bucket.bucketID] = (
               actionsByBucket[bucket.bucketID] || []
@@ -540,7 +568,7 @@ async function performDatabaseOperations(
             postID
           );
         } else if (name) {
-          const bucket = await dalBucket.getByName(name, posts[0].boardID);
+          const bucket = await dalBucket.getByName(name, socket.data.boardId);
           if (bucket) {
             actionsByBucket[bucket.bucketID] = (
               actionsByBucket[bucket.bucketID] || []
@@ -673,14 +701,9 @@ async function constructAndSendMessage(
   }
 }
 
-async function fetchAndFormatBuckets(posts: any[]): Promise<any[]> {
-  // Return empty list if no posts
-  if (!posts || posts.length === 0) {
-    return [];
-  }
-
+async function fetchAndFormatBuckets(boardId: string): Promise<any[]> {
   // Fetch ALL buckets for the board
-  const buckets = await dalBucket.getByBoardId(posts[0].boardID); // Assuming all posts belong to the same board
+  const buckets = await dalBucket.getByBoardId(boardId); // Assuming all posts belong to the same board
 
   // Create a list of buckets with their bucketIDs
   return buckets.map((bucket) => ({

diff --git a/backend/src/socket/events/ai.events.ts b/backend/src/socket/events/ai.events.ts
@@ -1,33 +1,35 @@
-// ai.events.ts
 import { Server, Socket } from 'socket.io';
 import { SocketEvent } from '../../constants';
 import { sendMessage } from '../../services/vertexAI';
 import { SocketPayload } from '../types/event.types'; // Import the type for SocketPayload
 
 // Define a type for the AI message data
 interface AiMessageData {
-  boardID: string;
   posts: any[]; // Or a more specific type for your posts
   prompt: string;
+  boardId: string;
+  userId: string;
 }
 
 class AiMessage {
   static type: SocketEvent = SocketEvent.AI_MESSAGE;
 
   static async handleEvent(
     data: SocketPayload<AiMessageData>
-  ): Promise<{ posts: any[]; prompt: string }> {
-    const { posts, prompt } = data.eventData;
+  ): Promise<{ posts: any[]; prompt: string; boardId: string; userId: string }> {
+    const { posts, prompt, boardId, userId } = data.eventData;
     // ... any necessary data processing or validation ...
-    return { posts, prompt };
+    return { posts, prompt, boardId, userId };
   }
 
   static async handleResult(
     io: Server,
     socket: Socket,
-    result: { posts: any[]; prompt: string; boardID: string }
+    result: { posts: any[]; prompt: string; boardId: string; userId: string }
   ): Promise<void> {
-    const { posts, prompt } = result;
+    const { posts, prompt, boardId, userId } = result;
+    socket.data.boardId = boardId;
+    socket.data.userId = userId;
     sendMessage(posts, prompt, socket);
   }
 }

diff --git a/frontend/src/app/components/create-workflow-modal/create-workflow-modal.component.html b/frontend/src/app/components/create-workflow-modal/create-workflow-modal.component.html
@@ -618,8 +618,11 @@ <h3 class="message-label">Assistant:</h3>
     <mat-label>Ask the AI Assistant</mat-label>
     <input matInput [(ngModel)]="aiPrompt" (keyup.enter)="askAI()" #aiInput /> 
   </mat-form-field>
-  <button mat-raised-button color="primary" (click)="askAI()" *ngIf="selected.value === 3">Send</button>
-
+  <div style="display: flex; gap: 10px;"> 
+    <button mat-raised-button color="primary" (click)="askAI()" *ngIf="selected.value === 3">Send</button>
+    <button mat-raised-button color="primary" (click)="downloadChatHistory()" *ngIf="selected.value === 3">Download Chat History</button>
+  </div>
+
   <button mat-button (click)="onNoClick()">Close</button>
   <button
     mat-button

diff --git a/frontend/src/app/components/create-workflow-modal/create-workflow-modal.component.ts b/frontend/src/app/components/create-workflow-modal/create-workflow-modal.component.ts
@@ -52,6 +52,7 @@ import { Subscription } from 'rxjs';
 import { UserService } from 'src/app/services/user.service';
 import User, { AuthUser, Role } from 'src/app/models/user';
 import { SocketEvent } from 'src/app/utils/constants';
+import { saveAs } from 'file-saver';
 
 interface ChatMessage {
   role: 'user' | 'assistant';
@@ -433,7 +434,12 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
       this.scrollToBottom();
 
       // 2. Send data and prompt to the backend via WebSocket
-      this.socketService.emit(SocketEvent.AI_MESSAGE, { posts, prompt });
+      this.socketService.emit(SocketEvent.AI_MESSAGE, {
+        posts,
+        prompt,
+        boardId: this.board.boardID,
+        userId: this.user.userID,
+      });
 
       // 3. Listen for WebSocket events
       this.aiResponseListener = this.socketService.listen(
@@ -474,8 +480,8 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
             }
 
             if (data.status === 'Completed' || data.status === 'Error') {
-              if (this.aiResponseListener) { 
-                this.aiResponseListener.unsubscribe(); 
+              if (this.aiResponseListener) {
+                this.aiResponseListener.unsubscribe();
               }
             }
 
@@ -489,8 +495,8 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
                 error,
             });
             this.stopWaitingForAIResponse();
-            if (this.aiResponseListener) { 
-              this.aiResponseListener.unsubscribe(); 
+            if (this.aiResponseListener) {
+              this.aiResponseListener.unsubscribe();
             }
           }
         }
@@ -499,6 +505,27 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
     this.isProcessingAIRequest = false;
   }
 
+  downloadChatHistory() {
+    const data = {
+      boardId: this.board.boardID,
+      userId: this.user.userID
+    };
+
+    this.http.post('chat-history', data, { 
+      responseType: 'blob' 
+    }).subscribe(
+      (response) => {
+        const blob = new Blob([response], { type: 'text/csv' });
+        saveAs(blob, 'chat_history.csv'); 
+        this.openSnackBar('Chat history downloaded successfully!');
+      },
+      (error) => {
+        console.error('Error downloading chat history:', error);
+        this.openSnackBar(`Error downloading chat history: ${error.message}`); 
+      }
+    );
+  }
+
   private escapeJsonResponse(response: string): string {
     if (!response) {
       return '';
@@ -858,8 +885,8 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
   }
 
   ngOnDestroy() {
-    if (this.aiResponseListener) {  
-      this.aiResponseListener.unsubscribe(); 
+    if (this.aiResponseListener) {
+      this.aiResponseListener.unsubscribe();
     }
   }
 }