vzcode 2.15.0 → 2.16.0

package/src/server/aiChatHandler/chatOperations.ts

@@ -1,11 +1,14 @@
 import { dateToTimestamp } from '@vizhub/viz-utils';
 import { randomId } from '../../randomId.js';
-import { ShareDBDoc } from '../../types.js';
+import {
+  ShareDBDoc,
+  StreamingEvent,
+  ExtendedVizContent,
+} from '../../types.js';
 import { diff } from '../../ot.js';
 import {
   VizChatId,
   VizContent,
-  VizFileId,
   VizFiles,
 } from '@vizhub/viz-types';
 
@@ -80,6 +83,8 @@ export const addUserMessage = (
   return userMessage;
 };
 
+const DEBUG = false;
+
 /**
  * Updates AI status in the chat
  */
@@ -88,6 +93,11 @@ export const updateAIStatus = (
   chatId: VizChatId,
   status: string,
 ) => {
+  DEBUG &&
+    console.log(
+      `ChatOperations: updateAIStatus called with status: "${status}" for chatId: ${chatId}`,
+    );
+
   const op = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
     chats: {
@@ -98,7 +108,17 @@ export const updateAIStatus = (
       },
     },
   });
+
+  DEBUG &&
+    console.log(
+      `ChatOperations: Submitting operation for status update:`,
+      op,
+    );
   shareDBDoc.submitOp(op);
+  DEBUG &&
+    console.log(
+      `ChatOperations: Status update operation submitted successfully`,
+    );
 };
 
 /**
@@ -405,69 +425,230 @@ export const createNewFile = (
   return newFileId;
 };
 
+// ============================================================================
+// Streaming Chat Operations
+// ============================================================================
+
 /**
- * Ensures a file exists, creating it if necessary
+ * Creates a streaming AI message with events array
  */
-export const ensureFileExists = (shareDBDoc, fileName) => {
-  let fileId = resolveFileId(fileName, shareDBDoc);
+export const createStreamingAIMessage = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+) => {
+  const aiMessage = {
+    id: `assistant-${Date.now()}`,
+    role: 'assistant',
+    content: '',
+    timestamp: dateToTimestamp(new Date()),
+    streamingEvents: [],
+    isProgressive: true,
+  };
 
-  if (!fileId) {
-    // File doesn't exist, create it
-    fileId = createNewFile(shareDBDoc, fileName);
-  }
+  const messageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        messages: [
+          ...shareDBDoc.data.chats[chatId].messages,
+          aiMessage,
+        ],
+        updatedAt: dateToTimestamp(new Date()),
+        isStreaming: true,
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
 
-  return fileId;
+  return aiMessage.id;
 };
 
 /**
- * Clears the content of a file
+ * Adds a streaming event to the most recent AI message
  */
-export const clearFileContent = (
-  shareDBDoc: ShareDBDoc<VizContent>,
-  fileId: VizFileId,
+export const addStreamingEvent = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+  event: StreamingEvent,
 ) => {
-  const currentFile = shareDBDoc.data.files[fileId];
+  DEBUG &&
+    console.log(
+      `ChatOperations: Adding streaming event:`,
+      event,
+    );
 
-  if (currentFile && currentFile.text) {
-    // Clear the file content
-    const newState = {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = [...chat.messages];
+  const lastMessageIndex = messages.length - 1;
+
+  if (
+    lastMessageIndex >= 0 &&
+    messages[lastMessageIndex].role === 'assistant'
+  ) {
+    const lastMessage = messages[lastMessageIndex] as any;
+    const updatedEvents = [
+      ...(lastMessage.streamingEvents || []),
+      event,
+    ];
+
+    messages[lastMessageIndex] = {
+      ...lastMessage,
+      streamingEvents: updatedEvents,
+    };
+
+    const messageOp = diff(shareDBDoc.data, {
       ...shareDBDoc.data,
-      files: {
-        ...shareDBDoc.data.files,
-        [fileId]: {
-          ...currentFile,
-          text: '',
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...chat,
+          messages,
+          updatedAt: dateToTimestamp(new Date()),
         },
       },
-    };
-
-    const op = diff(shareDBDoc.data, newState);
-    shareDBDoc.submitOp(op);
+    });
+    shareDBDoc.submitOp(messageOp);
   }
 };
 
 /**
- * Appends a line to a file using OT operations
+ * Updates streaming status for a chat
  */
-export const appendLineToFile = (
-  shareDBDoc: ShareDBDoc<VizContent>,
-  fileId: VizFileId,
-  line: string,
+export const updateStreamingStatus = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+  status: string,
+  isStreaming: boolean = true,
 ) => {
-  const currentFile = shareDBDoc.data.files[fileId];
-  const currentContent = currentFile?.text || '';
-  const newContent = currentContent + line + '\n';
+  DEBUG &&
+    console.log(
+      `ChatOperations: Updating streaming status: "${status}"`,
+    );
 
-  const newDocState = {
+  const op = diff(shareDBDoc.data, {
     ...shareDBDoc.data,
-    files: {
-      ...shareDBDoc.data.files,
-      [fileId]: {
-        ...currentFile,
-        text: newContent,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        currentStatus: status,
+        isStreaming,
+        updatedAt: dateToTimestamp(new Date()),
       },
     },
-  };
+  });
+  shareDBDoc.submitOp(op);
+};
 
-  shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
+/**
+ * Finalizes streaming message and clears streaming state
+ */
+export const finalizeStreamingMessage = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+) => {
+  DEBUG &&
+    console.log(
+      `ChatOperations: Finalizing streaming message`,
+    );
+
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = [...chat.messages];
+  const lastMessageIndex = messages.length - 1;
+
+  if (
+    lastMessageIndex >= 0 &&
+    messages[lastMessageIndex].role === 'assistant'
+  ) {
+    const lastMessage = messages[lastMessageIndex] as any;
+
+    messages[lastMessageIndex] = {
+      ...lastMessage,
+      isComplete: true,
+    };
+
+    const messageOp = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...chat,
+          messages,
+          currentStatus: 'Done',
+          isStreaming: false,
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+    shareDBDoc.submitOp(messageOp);
+  }
 };
+
+// /**
+//  * Ensures a file exists, creating it if necessary
+//  */
+// export const ensureFileExists = (shareDBDoc, fileName) => {
+//   let fileId = resolveFileId(fileName, shareDBDoc);
+
+//   if (!fileId) {
+//     // File doesn't exist, create it
+//     fileId = createNewFile(shareDBDoc, fileName);
+//   }
+
+//   return fileId;
+// };
+
+// /**
+//  * Clears the content of a file
+//  */
+// export const clearFileContent = (
+//   shareDBDoc: ShareDBDoc<VizContent>,
+//   fileId: VizFileId,
+// ) => {
+//   const currentFile = shareDBDoc.data.files[fileId];
+
+//   if (currentFile && currentFile.text) {
+//     // Clear the file content
+//     const newState = {
+//       ...shareDBDoc.data,
+//       files: {
+//         ...shareDBDoc.data.files,
+//         [fileId]: {
+//           ...currentFile,
+//           text: '',
+//         },
+//       },
+//     };
+
+//     const op = diff(shareDBDoc.data, newState);
+//     shareDBDoc.submitOp(op);
+//   }
+// };
+
+// /**
+//  * Appends a line to a file using OT operations
+//  */
+// export const appendLineToFile = (
+//   shareDBDoc: ShareDBDoc<VizContent>,
+//   fileId: VizFileId,
+//   line: string,
+// ) => {
+//   const currentFile = shareDBDoc.data.files[fileId];
+//   const currentContent = currentFile?.text || '';
+//   const newContent = currentContent + line + '\n';
+
+//   const newDocState = {
+//     ...shareDBDoc.data,
+//     files: {
+//       ...shareDBDoc.data.files,
+//       [fileId]: {
+//         ...currentFile,
+//         text: newContent,
+//       },
+//     },
+//   };
+
+//   shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
+// };

package/src/server/aiChatHandler/index.ts

@@ -3,14 +3,10 @@ import {
   ensureChatsExist,
   ensureChatExists,
   addUserMessage,
-  addDiffToAIMessage,
   setAIStatus,
 } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
-import {
-  performAIEditing,
-  performAIChat,
-} from './aiEditing.js';
+import { performAIEditing } from './aiEditing.js';
 import {
   handleError,
   handleBackgroundError,
@@ -27,18 +23,16 @@ export const handleAIChatMessage =
   ({
     shareDBDoc,
     onCreditDeduction,
-    getCurrentCommitId,
     model,
     aiRequestOptions,
   }: {
     shareDBDoc: ShareDBDoc<VizContent>;
     onCreditDeduction?: any;
-    getCurrentCommitId?: () => string | null;
     model?: string;
     aiRequestOptions?: any;
   }) =>
   async (req: any, res: any) => {
-    const { content, chatId, mode = 'edit' } = req.body;
+    const { content, chatId } = req.body;
 
     if (DEBUG) {
       console.log(
@@ -67,16 +61,11 @@ export const handleAIChatMessage =
       // Return success immediately - AI generation continues in background
       res.status(200).json('success');
 
-      // Set AI status to indicate generation is starting
-      setAIStatus(shareDBDoc, chatId, 'generating');
-
       // Continue AI processing in background (don't await)
       processAIRequestAsync({
         shareDBDoc,
         chatId,
         content,
-        mode,
-        getCurrentCommitId,
         model,
         aiRequestOptions,
         onCreditDeduction,
@@ -100,8 +89,6 @@ const processAIRequestAsync = async ({
   shareDBDoc,
   chatId,
   content,
-  mode,
-  getCurrentCommitId,
   model,
   aiRequestOptions,
   onCreditDeduction,
@@ -109,18 +96,11 @@ const processAIRequestAsync = async ({
   shareDBDoc: ShareDBDoc<VizContent>;
   chatId: string;
   content: string;
-  mode: string;
-  getCurrentCommitId?: () => string | null;
   model?: string;
   aiRequestOptions?: any;
   onCreditDeduction?: any;
 }) => {
   try {
-    // Capture the current commit ID before making changes (for VizHub integration)
-    const beforeCommitId = getCurrentCommitId
-      ? getCurrentCommitId()
-      : null;
-
     // Create LLM function for streaming
     const llmFunction = createLLMFunction({
       shareDBDoc,
@@ -135,32 +115,12 @@ const processAIRequestAsync = async ({
     const runCode = createRunCodeFunction(submitOperation);
 
     // Perform AI editing or chat based on mode
-    const editResult =
-      mode === 'ask'
-        ? await performAIChat({
-            prompt: content,
-            shareDBDoc,
-            llmFunction,
-          })
-        : await performAIEditing({
-            prompt: content,
-            shareDBDoc,
-            llmFunction,
-            runCode,
-          });
-
-    // Add diff data to the AI message if there are changes
-    if (
-      editResult.diffData &&
-      Object.keys(editResult.diffData).length > 0
-    ) {
-      addDiffToAIMessage(
-        shareDBDoc,
-        chatId,
-        editResult.diffData,
-        beforeCommitId, // Pass the commit ID before AI changes (for VizHub integration)
-      );
-    }
+    const editResult = await performAIEditing({
+      prompt: content,
+      shareDBDoc,
+      llmFunction,
+      runCode,
+    });
 
     // Handle credit deduction if callback is provided
     if (onCreditDeduction && editResult.generationId) {