vzcode 2.18.0 → 2.22.0

package/src/llm-streaming-server/chatOperations.ts

@@ -0,0 +1,655 @@
+import { dateToTimestamp } from '@vizhub/viz-utils';
+import { randomId } from '../randomId.js';
+import {
+  ShareDBDoc,
+  StreamingEvent,
+  ExtendedVizContent,
+} from '../types.js';
+import { diff } from '../ot.js';
+import {
+  VizChatId,
+  VizContent,
+  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
+ * Clears old messages to reflect that each prompt is a self-contained code transformation
+ */
+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: [userMessage], // Replace old messages with just the new user message
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(userMessageOp);
+
+  return userMessage;
+};
+
+const DEBUG = false;
+
+/**
+ * Updates AI status in the chat
+ */
+export const updateAIStatus = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  status: string,
+) => {
+  DEBUG &&
+    console.log(
+      `ChatOperations: updateAIStatus called with status: "${status}" for chatId: ${chatId}`,
+    );
+
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiStatus: status,
+      },
+    },
+  });
+
+  DEBUG &&
+    console.log(
+      `ChatOperations: Submitting operation for status update:`,
+      op,
+    );
+  shareDBDoc.submitOp(op);
+  DEBUG &&
+    console.log(
+      `ChatOperations: Status update operation submitted successfully`,
+    );
+};
+
+/**
+ * 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,
+  });
+  DEBUG && console.log('updateFiles op:');
+  DEBUG && console.log(JSON.stringify(filesOp, null, 2));
+  shareDBDoc.submitOp(filesOp);
+  return 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;
+};
+
+// ============================================================================
+// Streaming Chat Operations
+// ============================================================================
+
+/**
+ * Creates a streaming AI message with events array
+ */
+export const createStreamingAIMessage = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+) => {
+  const aiMessage = {
+    id: `assistant-${Date.now()}`,
+    role: 'assistant',
+    content: '',
+    timestamp: dateToTimestamp(new Date()),
+    streamingEvents: [],
+    isProgressive: true,
+  };
+
+  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 aiMessage.id;
+};
+
+/**
+ * Adds a streaming event to the most recent AI message
+ */
+export const addStreamingEvent = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+  event: StreamingEvent,
+) => {
+  DEBUG &&
+    console.log(
+      `ChatOperations: Adding streaming event:`,
+      event,
+    );
+
+  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,
+      chats: {
+        ...shareDBDoc.data.chats,
+        [chatId]: {
+          ...chat,
+          messages,
+          updatedAt: dateToTimestamp(new Date()),
+        },
+      },
+    });
+    shareDBDoc.submitOp(messageOp);
+  }
+};
+
+/**
+ * Updates streaming status for a chat
+ */
+export const updateStreamingStatus = (
+  shareDBDoc: ShareDBDoc<ExtendedVizContent>,
+  chatId: VizChatId,
+  status: string,
+  isStreaming: boolean = true,
+) => {
+  DEBUG &&
+    console.log(
+      `ChatOperations: Updating streaming status: "${status}"`,
+    );
+
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        currentStatus: status,
+        isStreaming,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * 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/llm-streaming-server/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);
+};