Merge pull request #626 from encorelab/625-feature-add-ai-assistant-c…

…hat-log-downloads

Added download chat history button for AI Assistant

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;

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;

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;

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'));

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';
@@ -341,6 +342,14 @@ function parseJsonResponse(response: string): any {
   }
 }
 
+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,
@@ -350,6 +359,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);
 
@@ -364,7 +384,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(
@@ -419,6 +439,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 = `Completed with invalid formatting: ${partialResponse}`;
             socket.emit(SocketEvent.AI_RESPONSE, {
@@ -476,7 +505,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));
 
@@ -486,7 +515,7 @@ async function performDatabaseOperations(
 
       const newBucket: BucketModel = {
         bucketID: generateUniqueID(),
-        boardID: posts[0].boardID,
+        boardID: socket.data.boardId,
         name: bucketName,
         posts: [],
       };
@@ -505,7 +534,7 @@ async function performDatabaseOperations(
 
       const newBucket: BucketModel = {
         bucketID: generateUniqueID(),
-        boardID: posts[0].boardID,
+        boardID: socket.data.boardId,
         name: name,
         posts: [],
       };
@@ -545,7 +574,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] || []
@@ -591,7 +620,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] || []
@@ -724,14 +753,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) => ({

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);
   }
 }

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