vzcode 1.32.0 → 1.34.0

package/src/client/vzReducer/index.ts

@@ -30,6 +30,10 @@ import {
   setSearchFocusedIndexReducer,
   toggleSearchFocusedReducer,
 } from './searchReducer';
+import {
+  setIsAIChatOpenReducer,
+  toggleAIChatFocusedReducer,
+} from './aiChatReducer';
 import { toggleAutoFollowReducer } from './toggleAutoFollowReducer';
 import { updatePresenceIndicatorReducer } from './updatePresenceIndicatorReducer';
 import { splitCurrentPaneReducer } from './splitCurrentPaneReducer';
@@ -52,6 +56,12 @@ export type VZState = {
   // True to show the search instead of files
   isSearchOpen: boolean;
 
+  // True to show the AI chat instead of files
+  isAIChatOpen: boolean;
+
+  // True if the AI chat input should focus on the next render.
+  aiChatFocused: boolean;
+
   // True to show the settings modal.
   isSettingsOpen: boolean;
 
@@ -115,6 +125,14 @@ export type VZAction =
   //  * Sets whether the search tab is open.
   | { type: 'set_is_search_open'; value: boolean }
 
+  // `set_is_ai_chat_open`
+  //  * Sets whether the AI chat tab is open.
+  | { type: 'set_is_ai_chat_open'; value: boolean }
+
+  // `toggle_ai_chat_focused`
+  // * Toggles focused variable to trigger AI chat input focus
+  | { type: 'toggle_ai_chat_focused' }
+
   // `set_search`
   //  * Sets the current search pattern
   | { type: 'set_search'; value: string }
@@ -192,6 +210,8 @@ const reducers = [
   setSearchFocusedIndexReducer,
   toggleSearchFocusedReducer,
   setIsSearchOpenReducer,
+  setIsAIChatOpenReducer,
+  toggleAIChatFocusedReducer,
   setIsSettingsOpenReducer,
   setIsDocOpenReducer,
   editorNoLongerWantsFocusReducer,

package/src/ot.js

@@ -1,5 +1,13 @@
+// Operational Transformation (OT) utilities.
 // This file is the central point where the OT types are imported.
 // Localized to one file so it's easy to change it in future.
 import OTJSON1Presence from 'sharedb-client-browser/dist/ot-json1-presence-umd.cjs';
-const { json1Presence, textUnicode } = OTJSON1Presence;
-export { json1Presence, textUnicode };
+
+export const { json1Presence, textUnicode } =
+  OTJSON1Presence;
+
+// The OT type used throughout the codebase.
+export const otType = json1Presence.type;
+
+// Applies an OT op to an object.
+export const apply = otType.apply;

package/src/server/aiChatHandler/aiEditing.js

@@ -0,0 +1,42 @@
+import { performAiEdit } from 'editcodewithai';
+import {
+  clearAIScratchpadAndStatus,
+  updateFiles,
+  setIsInteracting,
+} from './chatOperations.js';
+
+/**
+ * Performs AI editing operations on the files
+ */
+export const performAIEditing = async (
+  shareDBDoc,
+  chatId,
+  content,
+  files,
+  llmFunction,
+) => {
+  const editResult = await performAiEdit({
+    prompt: content,
+    files,
+    llmFunction,
+    apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
+  });
+
+  // Clear the scratchpad and update status
+  clearAIScratchpadAndStatus(
+    shareDBDoc,
+    chatId,
+    'Done editing with AI.',
+  );
+
+  // Apply AI edits to files
+  updateFiles(shareDBDoc, editResult.changedFiles);
+
+  // Wait for propagation
+  await new Promise((resolve) => setTimeout(resolve, 100));
+
+  // Unset isInteracting
+  setIsInteracting(shareDBDoc, false);
+
+  return editResult;
+};

package/src/server/aiChatHandler/chatOperations.js

@@ -0,0 +1,195 @@
+import { dateToTimestamp } from '@vizhub/viz-utils';
+import { diff } from '../../client/diff.js';
+
+/**
+ * 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, chatId) => {
+  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,
+  chatId,
+  content,
+) => {
+  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,
+  chatId,
+  status,
+) => {
+  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,
+  chatId,
+  content,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: content,
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Clears AI scratchpad and updates status
+ */
+export const clearAIScratchpadAndStatus = (
+  shareDBDoc,
+  chatId,
+  status,
+) => {
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...shareDBDoc.data.chats[chatId],
+        aiScratchpad: undefined,
+        aiStatus: status,
+      },
+    },
+  });
+  shareDBDoc.submitOp(op);
+};
+
+/**
+ * Adds an AI response message to the chat
+ */
+export const addAIMessage = (
+  shareDBDoc,
+  chatId,
+  content,
+) => {
+  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, files) => {
+  const filesOp = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    files: files,
+    isInteracting: true,
+  });
+  shareDBDoc.submitOp(filesOp);
+};
+
+/**
+ * Sets isInteracting flag
+ */
+export const setIsInteracting = (
+  shareDBDoc,
+  isInteracting,
+) => {
+  const interactingOp = diff(
+    { isInteracting: !isInteracting },
+    { isInteracting: isInteracting },
+  );
+  shareDBDoc.submitOp(interactingOp);
+};

package/src/server/aiChatHandler/errorHandling.js

@@ -0,0 +1,53 @@
+import { dateToTimestamp } from '@vizhub/viz-utils';
+import { diff } from '../../client/diff.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,
+    );
+  }
+
+  res.status(500).json({
+    error: 'Internal server error',
+    message: error.message,
+  });
+};

package/src/server/aiChatHandler/index.js

@@ -0,0 +1,69 @@
+import { validateRequest } from './validation.js';
+import {
+  ensureChatsExist,
+  ensureChatExists,
+  addUserMessage,
+  addAIMessage,
+} from './chatOperations.js';
+import { createLLMFunction } from './llmStreaming.js';
+import { performAIEditing } from './aiEditing.js';
+import { handleError } from './errorHandling.js';
+
+const debug = false;
+
+export const handleAIChatMessage =
+  (shareDBDoc) => async (req, res) => {
+    const { content, chatId } = req.body;
+
+    if (debug) {
+      console.log(
+        '[handleAIChatMessage] content:',
+        content,
+        'chatId:',
+        chatId,
+      );
+    }
+
+    // Validate request
+    if (!validateRequest(req, res)) {
+      return;
+    }
+
+    try {
+      // Get existing files from ShareDB doc
+      const files = shareDBDoc.data.files;
+
+      // Ensure chats structure exists
+      ensureChatsExist(shareDBDoc);
+      ensureChatExists(shareDBDoc, chatId);
+
+      // Add user message to chat
+      addUserMessage(shareDBDoc, chatId, content);
+
+      // Create LLM function for streaming
+      const llmFunction = createLLMFunction(
+        shareDBDoc,
+        chatId,
+      );
+
+      // Perform AI editing
+      const editResult = await performAIEditing(
+        shareDBDoc,
+        chatId,
+        content,
+        files,
+        llmFunction,
+      );
+
+      // Add AI response message
+      const aiResponse = addAIMessage(
+        shareDBDoc,
+        chatId,
+        editResult.content,
+      );
+
+      res.status(200).json(aiResponse);
+    } catch (error) {
+      handleError(shareDBDoc, chatId, error, res);
+    }
+  };

package/src/server/aiChatHandler/llmStreaming.js

@@ -0,0 +1,105 @@
+import { StreamingMarkdownParser } from 'llm-code-format';
+import { ChatOpenAI } from '@langchain/openai';
+import {
+  updateAIStatus,
+  updateAIScratchpad,
+} from './chatOperations.js';
+
+const debug = false;
+
+/**
+ * Creates and configures the LLM function for streaming
+ */
+export const createLLMFunction = (shareDBDoc, chatId) => {
+  return async (fullPrompt) => {
+    const chatModel = new ChatOpenAI({
+      modelName:
+        process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
+        'gpt-4o-mini',
+      configuration: {
+        apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
+        baseURL: process.env.VIZHUB_EDIT_WITH_AI_BASE_URL,
+        defaultHeaders: {
+          'HTTP-Referer': 'https://vizhub.com',
+          'X-Title': 'VizHub',
+        },
+      },
+      streaming: true,
+    });
+
+    let fullContent = '';
+    let generationId = '';
+
+    // Track expected state to avoid race conditions
+    let expectedAiScratchpad =
+      shareDBDoc.data.chats[chatId]?.aiScratchpad || '';
+
+    // Throttle updates to avoid "op too long" errors
+    let lastUpdateTime = 0;
+    const UPDATE_THROTTLE_MS = 100;
+
+    // Define callbacks for streaming parser
+    const callbacks = {
+      onFileNameChange: (fileName, format) => {
+        debug &&
+          console.log(
+            `File changed to: ${fileName} (${format})`,
+          );
+        updateAIStatus(
+          shareDBDoc,
+          chatId,
+          'Editing ' + fileName,
+        );
+      },
+      onCodeLine: (line) => {
+        debug && console.log(`Code line: ${line}`);
+      },
+      onNonCodeLine: (line) => {
+        debug && console.log(`Comment/text: ${line}`);
+      },
+    };
+
+    const parser = new StreamingMarkdownParser(callbacks);
+
+    // Stream the response
+    const stream = await chatModel.stream(fullPrompt);
+    for await (const chunk of stream) {
+      if (chunk.content) {
+        const chunkContent = String(chunk.content);
+        fullContent += chunkContent;
+
+        // Throttle updates
+        const now = Date.now();
+        const shouldUpdate =
+          now - lastUpdateTime >= UPDATE_THROTTLE_MS;
+
+        if (shouldUpdate) {
+          updateAIScratchpad(
+            shareDBDoc,
+            chatId,
+            fullContent,
+          );
+          expectedAiScratchpad = fullContent;
+          lastUpdateTime = now;
+        }
+
+        parser.processChunk(chunkContent);
+      }
+
+      if (!generationId && chunk.lc_kwargs?.id) {
+        generationId = chunk.lc_kwargs.id;
+      }
+    }
+    parser.flushRemaining();
+
+    // Submit final update if there's pending content
+    if (expectedAiScratchpad !== fullContent) {
+      updateAIScratchpad(shareDBDoc, chatId, fullContent);
+    }
+
+    return {
+      content: fullContent,
+      generationId: generationId,
+    };
+  };
+};

package/src/server/aiChatHandler/validation.js

@@ -0,0 +1,24 @@
+/**
+ * Validates incoming request data for AI chat messages
+ */
+export const validateRequest = (req, res) => {
+  const { content, chatId } = req.body;
+
+  if (!content || typeof content !== 'string') {
+    res.status(400).json({
+      error:
+        'Invalid request: content is required and must be a string',
+    });
+    return false;
+  }
+
+  if (!chatId || typeof chatId !== 'string') {
+    res.status(400).json({
+      error:
+        'Invalid request: chatId is required and must be a string',
+    });
+    return false;
+  }
+
+  return true;
+};

package/src/server/handleAIChatMessage.js

@@ -0,0 +1,2 @@
+// Re-export the refactored handleAIChatMessage from the new modular structure
+export { handleAIChatMessage } from './aiChatHandler/index.js';

package/src/server/index.js

@@ -14,6 +14,7 @@ import { json1Presence } from '../ot.js';
 import { computeInitialDocument } from './computeInitialDocument.js';
 import { handleAIAssist } from './handleAIAssist.js';
 import { handleAICopilot } from './handleAICopilot.js';
+import { handleAIChatMessage } from './handleAIChatMessage.js';
 import { isDirectory } from './isDirectory.js';
 import { createToken } from './livekit.js';
 import './setupEnv.js';
@@ -112,6 +113,13 @@ app.post(
   handleAICopilot(),
 );
 
+// Handle AI Chat Message requests.
+app.post(
+  '/ai-chat-message',
+  bodyParser.json(),
+  handleAIChatMessage(shareDBDoc),
+);
+
 // Livekit Token Generator
 app.get('/livekit-token', async (req, res) => {
   const { room, username } = req.query;