vzcode 2.3.0 → 2.5.0

package/src/server/aiChatHandler/chatOperations.ts

@@ -0,0 +1,476 @@
+import {
+  dateToTimestamp,
+  generateRunId,
+} from '@vizhub/viz-utils';
+import { randomId } from '../../randomId.js';
+import { ShareDBDoc } from '../../types.js';
+import { diff } from '../../ot.js';
+import {
+  VizChatId,
+  VizContent,
+  VizFileId,
+  VizFiles,
+} from '@vizhub/viz-types';
+
+/**
+ * Ensures the chats object exists in the ShareDB document
+ */
+export const ensureChatsExist = (shareDBDoc) => {
+  if (!shareDBDoc.data.chats) {
+    const op = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {},
+    });
+    shareDBDoc.submitOp(op);
+  }
+};
+
+/**
+ * Ensures a specific chat exists in the ShareDB document
+ */
+export const ensureChatExists = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
+  if (!shareDBDoc.data.chats[chatId]) {
+    const op = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          id: chatId,
+          messages: [],
+          createdAt: dateToTimestamp(new Date()),
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+    shareDBDoc.submitOp(op);
+  }
+};
+
+/**
+ * Adds a user message to the chat
+ */
+export const addUserMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content: string,
+) => {
+  const userMessage = {
+    id: `user-${Date.now()}`,
+    role: 'user',
+    content: content,
+    timestamp: dateToTimestamp(new Date()),
+  };
+
+  const userMessageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        messages: [
+          ...shareDBDoc.data.chats[chatId].messages,
+          userMessage,
+        ],
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(userMessageOp);
+
+  return userMessage;
+};
+
+/**
+ * Updates AI status in the chat
+ */
+export const updateAIStatus = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiStatus: status,
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Updates AI scratchpad content
+ */
+export const updateAIScratchpad = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content: string,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: content,
+      },
+    },
+  });
+
+  // op is `null` if there are no changes
+  // This can happen if the content is the same as before
+  // In that case, we don't need to submit an operation
+  if (op) {
+    shareDBDoc.submitOp(op);
+  }
+};
+
+/**
+ * Clears AI scratchpad and updates status
+ */
+export const clearAIScratchpadAndStatus = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: undefined,
+        aiStatus: status,
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Creates an initial empty AI message for streaming
+ */
+export const createAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
+  const aiMessage = {
+    id: `assistant-${Date.now()}`,
+    role: 'assistant',
+    content: '',
+    timestamp: dateToTimestamp(new Date()),
+  };
+
+  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()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
+
+  return aiMessage.id;
+};
+
+/**
+ * Updates the content of an AI message during streaming
+ */
+export const updateAIMessageContent = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  messageId: string,
+  content: string,
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messageIndex = chat.messages.findIndex(
+    (msg) => msg.id === messageId,
+  );
+
+  if (messageIndex === -1) {
+    console.warn(
+      `AI message with id ${messageId} not found`,
+    );
+    return;
+  }
+
+  const updatedMessages = [...chat.messages];
+  updatedMessages[messageIndex] = {
+    ...updatedMessages[messageIndex],
+    content,
+  };
+
+  const messageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...chat,
+        messages: updatedMessages,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
+};
+
+/**
+ * Sets the AI status for a chat
+ */
+export const setAIStatus = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string | undefined,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiStatus: status,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Finalizes an AI message by clearing temporary fields
+ */
+export const finalizeAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: undefined,
+        aiStatus: undefined,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Adds diff data to the most recent AI message
+ */
+export const addDiffToAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  diffData: any,
+  beforeCommitId?: string, // Commit ID before AI changes for VizHub integration
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = [...chat.messages];
+
+  // Find the most recent AI message
+  const lastAIMessageIndex = messages.length - 1;
+  if (
+    lastAIMessageIndex >= 0 &&
+    messages[lastAIMessageIndex].role === 'assistant'
+  ) {
+    const newMessage = {
+      ...messages[lastAIMessageIndex],
+      diffData,
+      ...(beforeCommitId && { beforeCommitId }), // Add beforeCommitId for VizHub integration
+    };
+    // Use type assertion to extend the message with diffData
+    (messages[lastAIMessageIndex] as any) = newMessage;
+
+    const messageOp = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...chat,
+          messages,
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+
+    shareDBDoc.submitOp(messageOp);
+  }
+};
+
+/**
+ * Adds an AI response message to the chat (legacy function, kept for compatibility)
+ */
+export const addAIMessage = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  content?: string,
+) => {
+  const aiResponse = {
+    id: Date.now() + 1,
+    role: 'assistant',
+    content: content || 'AI edit completed successfully.',
+    timestamp: dateToTimestamp(new Date()),
+  };
+
+  const messageOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        messages: [
+          ...shareDBDoc.data.chats[chatId].messages,
+          aiResponse,
+        ],
+        updatedAt: dateToTimestamp(new Date()),
+        aiStatus: undefined,
+      },
+    },
+  });
+  shareDBDoc.submitOp(messageOp);
+
+  return aiResponse;
+};
+
+/**
+ * Updates files in the ShareDB document
+ */
+export const updateFiles = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  files: VizFiles,
+) => {
+  const filesOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    files,
+  });
+  shareDBDoc.submitOp(filesOp);
+};
+
+/**
+ * Finds a file ID by searching for a matching file name
+ */
+export const resolveFileId = (
+  fileName: string,
+  shareDBDoc: any,
+) => {
+  const files = shareDBDoc.data.files;
+
+  // Search through all files to find matching name
+  for (const [fileId, file] of Object.entries(files)) {
+    if ((file as any).name === fileName) {
+      return fileId;
+    }
+  }
+
+  // If file doesn't exist, return null
+  return null;
+};
+
+/**
+ * Creates a new file with a random ID
+ */
+export const createNewFile = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  fileName: string,
+) => {
+  // Generate a new random file ID
+  const newFileId = randomId();
+
+  const newState = {
+    ...shareDBDoc.data,
+    files: {
+      ...shareDBDoc.data.files,
+      [newFileId]: {
+        name: fileName,
+        text: '',
+      },
+    },
+  };
+
+  const op = diff(shareDBDoc.data, newState);
+  shareDBDoc.submitOp(op);
+  return newFileId;
+};
+
+/**
+ * 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/errorHandling.ts

@@ -0,0 +1,66 @@
+import { dateToTimestamp } from '@vizhub/viz-utils';
+import { diff } from '../../ot.js';
+
+/**
+ * Handles errors by adding an error message to the chat and clearing AI state
+ */
+export const handleError = (
+  shareDBDoc,
+  chatId,
+  error,
+  res,
+) => {
+  console.error('[handleAIChatMessage] error:', error);
+
+  // Clear scratchpad and add error message on error
+  try {
+    const errorResponse = {
+      id: `error-${Date.now()}`,
+      role: 'assistant',
+      content:
+        'Sorry, I encountered an error while processing your message. Please try again.',
+      timestamp: dateToTimestamp(new Date()),
+    };
+
+    const errorOp = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...shareDBDoc.data.chats[chatId],
+          messages: [
+            ...shareDBDoc.data.chats[chatId].messages,
+            errorResponse,
+          ],
+          aiScratchpad: undefined,
+          aiStatus: undefined,
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+    shareDBDoc.submitOp(errorOp);
+  } catch (opError) {
+    console.error(
+      '[handleAIChatMessage] error handling error state:',
+      opError,
+    );
+  }
+
+  if (res) {
+    res.status(500).json({
+      error: 'Internal server error',
+      message: error.message,
+    });
+  }
+};
+
+/**
+ * Handles errors in background processing (without HTTP response)
+ */
+export const handleBackgroundError = (
+  shareDBDoc,
+  chatId,
+  error,
+) => {
+  handleError(shareDBDoc, chatId, error, null);
+};