vzcode 1.44.0 → 1.46.0

package/src/server/aiChatHandler/llmStreaming.ts

@@ -3,7 +3,9 @@ import {
   StreamingMarkdownParser,
 } from 'llm-code-format';
 import { ChatOpenAI } from '@langchain/openai';
+import OpenAI from 'openai';
 import fs from 'fs';
+import { generateRunId } from '@vizhub/viz-utils';
 import {
   updateAIStatus,
   createAIMessage,
@@ -13,8 +15,12 @@ import {
   clearFileContent,
   appendLineToFile,
   updateFiles,
+  updateAIScratchpad,
 } from './chatOperations.js';
 import { mergeFileChanges } from 'editcodewithai';
+import { diff } from '../../ot.js';
+import { VizChatId, VizContent } from '@vizhub/viz-types';
+import { ShareDBDoc } from '../../types.js';
 
 const DEBUG = false;
 
@@ -34,28 +40,30 @@ const DEBUG = false;
 const enableStreamingEditing = false;
 
 /**
- * Creates and configures the LLM function for streaming
+ * Creates and configures the LLM function for streaming with reasoning tokens
  */
 export const createLLMFunction = ({
   shareDBDoc,
   createVizBotLocalPresence,
   chatId,
+}: {
+  shareDBDoc: ShareDBDoc<VizContent>;
+  createVizBotLocalPresence: () => any;
+  chatId: VizChatId;
 }) => {
   return async (fullPrompt: string) => {
     const localPresence = createVizBotLocalPresence();
-    const chatModel = new ChatOpenAI({
-      modelName:
-        process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
-        'anthropic/claude-sonnet-4',
-      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',
-        },
+
+    // Create OpenRouter client for reasoning token support
+    const openRouterClient = new OpenAI({
+      apiKey: process.env.VIZHUB_EDIT_WITH_AI_API_KEY,
+      baseURL:
+        process.env.VIZHUB_EDIT_WITH_AI_BASE_URL ||
+        'https://openrouter.ai/api/v1',
+      defaultHeaders: {
+        'HTTP-Referer': 'https://vizhub.com',
+        'X-Title': 'VizHub',
       },
-      streaming: true,
     });
 
     let fullContent = '';
@@ -177,12 +185,58 @@ export const createLLMFunction = ({
     const parser = new StreamingMarkdownParser(callbacks);
 
     const chunks = [];
+    let reasoningContent = '';
+
+    // Stream the response with reasoning tokens
+    const modelName =
+      process.env.VIZHUB_EDIT_WITH_AI_MODEL_NAME ||
+      'anthropic/claude-3.5-sonnet';
+    const stream = await (
+      openRouterClient.chat.completions.create as any
+    )({
+      model: modelName,
+      messages: [{ role: 'user', content: fullPrompt }],
+      max_tokens: 8192,
+      reasoning: {
+        effort: 'medium',
+        exclude: false,
+      },
+      usage: { include: true },
+      stream: true,
+    });
+
+    let reasoningStarted = false;
+    let contentStarted = false;
 
-    // Stream the response
-    const stream = await chatModel.stream(fullPrompt);
     for await (const chunk of stream) {
-      if (chunk && chunk.content) {
-        const chunkContent = String(chunk.content);
+      const delta = chunk.choices[0]?.delta as any; // Type assertion for OpenRouter-specific reasoning fields
+
+      if (delta?.reasoning) {
+        // Handle reasoning tokens (thinking)
+        if (!reasoningStarted) {
+          reasoningStarted = true;
+          updateAIStatus(shareDBDoc, chatId, 'Thinking...');
+        }
+        reasoningContent += delta.reasoning;
+        updateAIScratchpad(
+          shareDBDoc,
+          chatId,
+          reasoningContent,
+        );
+      } else if (delta?.content) {
+        // Handle regular content tokens
+        if (reasoningStarted && !contentStarted) {
+          // Clear reasoning when content starts
+          contentStarted = true;
+          updateAIScratchpad(shareDBDoc, chatId, '');
+          updateAIStatus(
+            shareDBDoc,
+            chatId,
+            'Generating response...',
+          );
+        }
+
+        const chunkContent = delta.content;
         chunks.push(chunkContent);
 
         if (enableStreamingEditing) {
@@ -196,14 +250,20 @@ export const createLLMFunction = ({
             fullContent,
           );
         }
+      } else if (chunk.usage) {
+        // Handle usage information
+        DEBUG && console.log('Usage:', chunk.usage);
       }
 
-      if (!generationId && chunk.lc_kwargs?.id) {
-        generationId = chunk.lc_kwargs.id;
+      if (!generationId && chunk.id) {
+        generationId = chunk.id;
       }
     }
     await parser.flushRemaining();
     reportFileEdited();
+
+    // Final cleanup - clear scratchpad and set final status
+    updateAIScratchpad(shareDBDoc, chatId, '');
     updateAIStatus(shareDBDoc, chatId, 'Done editing.');
     updateAIMessageContent(
       shareDBDoc,
@@ -242,6 +302,28 @@ export const createLLMFunction = ({
       );
     }
 
+    // Generate a new runId to trigger a run when AI finishes editing
+    // This will trigger a re-run without hot reloading
+    const newRunId = generateRunId();
+    const runIdOp = diff(shareDBDoc.data, {
+      ...shareDBDoc.data,
+      runId: newRunId,
+    });
+    shareDBDoc.submitOp(runIdOp, (error) => {
+      if (error) {
+        console.warn(
+          'Error setting runId after AI editing:',
+          error,
+        );
+      } else {
+        DEBUG &&
+          console.log(
+            'Set new runId after AI editing:',
+            newRunId,
+          );
+      }
+    });
+
     // Write chunks file for debugging
     if (DEBUG) {
       const chunksFileJSONpath = `./ai-chunks-${chatId}.json`;

package/src/server/computeInitialDocument.ts

@@ -10,6 +10,15 @@ import createIgnore from 'ignore';
 import { ignoreFilePattern, baseIgnore } from './config.js';
 import { isDirectory } from './isDirectory.js';
 
+// Import the image file utility
+const isImageFile = (fileName) => {
+  return (
+    fileName.match(
+      /\.(png|jpg|jpeg|gif|bmp|svg|webp)$/i,
+    ) !== null
+  );
+};
+
 /**
  * @param {string} fullPath - absolut path of the workspace root
  * @param {string} currentDirectoryPath - path where the ignore file is found, relative to fullPath
@@ -180,13 +189,23 @@ export const computeInitialDocument = ({ fullPath }) => {
 
   files.forEach((file) => {
     const id = randomId();
+    let text = null;
+
+    if (!isDirectory(file)) {
+      const filePath = path.join(fullPath, file);
+
+      if (isImageFile(file)) {
+        // Read image files as binary and convert to base64
+        const buffer = fs.readFileSync(filePath);
+        text = buffer.toString('base64');
+      } else {
+        // Read non-image files as UTF-8 text
+        text = fs.readFileSync(filePath, 'utf-8');
+      }
+    }
+
     initialDocument.files[id] = {
-      text: isDirectory(file)
-        ? null
-        : fs.readFileSync(
-            path.join(fullPath, file),
-            'utf-8',
-          ),
+      text,
       name: file,
     };
   });

package/src/server/generateAIResponse.ts

@@ -150,6 +150,10 @@ export const generateAIResponse = async ({
     // model: 'gpt-3.5-turbo',
     model: 'gpt-4o',
     messages,
+    reasoning: {
+      effort: 'medium',
+      exclude: false,
+    } as any, // Type assertion for OpenRouter-specific reasoning parameter
     stream: true,
   });
 

package/src/server/index.ts

@@ -14,10 +14,20 @@ import { computeInitialDocument } from './computeInitialDocument.js';
 import { handleAIAssist } from './handleAIAssist.js';
 import { handleAICopilot } from './handleAICopilot.js';
 import { handleAIChatMessage } from './handleAIChatMessage.js';
+import { undoAIEdit } from './aiChatHandler/chatOperations.js';
 import { isDirectory } from './isDirectory.js';
 import { createToken } from './livekit.js';
 import './setupEnv.js';
 
+// Import the image file utility
+const isImageFile = (fileName) => {
+  return (
+    fileName.match(
+      /\.(png|jpg|jpeg|gif|bmp|svg|webp)$/i,
+    ) !== null
+  );
+};
+
 // The time in milliseconds by which auto-saving is debounced.
 const autoSaveDebounceTimeMS = 800;
 
@@ -132,6 +142,57 @@ app.post(
   }),
 );
 
+// Handle AI Chat Undo requests.
+app.post(
+  '/ai-chat-undo',
+  bodyParser.json(),
+  async (req, res) => {
+    const { chatId, messageId } = req.body;
+
+    if (!chatId || !messageId) {
+      return res
+        .status(400)
+        .json({ error: 'Missing chatId or messageId' });
+    }
+
+    try {
+      const chat = shareDBDoc.data.chats?.[chatId];
+      if (!chat) {
+        return res
+          .status(404)
+          .json({ error: 'Chat not found' });
+      }
+
+      const message = chat.messages.find(
+        (msg) => msg.id === messageId,
+      );
+      if (
+        !message ||
+        message.role !== 'assistant' ||
+        !message.beforeFiles
+      ) {
+        return res
+          .status(400)
+          .json({ error: 'Invalid message for undo' });
+      }
+
+      undoAIEdit(
+        shareDBDoc,
+        chatId,
+        messageId,
+        message.beforeFiles,
+      );
+
+      res.status(200).json({ success: true });
+    } catch (error) {
+      console.error('Error undoing AI edit:', error);
+      res
+        .status(500)
+        .json({ error: 'Failed to undo edit' });
+    }
+  },
+);
+
 // Livekit Token Generator
 app.get('/livekit-token', async (req, res) => {
   const { room, username } = req.query;
@@ -161,7 +222,17 @@ const save = () => {
     if (previous && current) {
       // Handle changing of text content.
       if (previous.text !== current.text) {
-        fs.writeFileSync(current.name, current.text);
+        if (isImageFile(current.name)) {
+          // Write image files as binary from base64
+          const buffer = Buffer.from(
+            current.text,
+            'base64',
+          );
+          fs.writeFileSync(current.name, buffer);
+        } else {
+          // Write non-image files as text
+          fs.writeFileSync(current.name, current.text);
+        }
       }
 
       // // Handle renaming files.
@@ -256,7 +327,17 @@ const save = () => {
     if (!previous && current) {
       //File Creation
       if (!isDirectory(current.name)) {
-        fs.writeFileSync(current.name, current.text);
+        if (isImageFile(current.name)) {
+          // Write image files as binary from base64
+          const buffer = Buffer.from(
+            current.text,
+            'base64',
+          );
+          fs.writeFileSync(current.name, buffer);
+        } else {
+          // Write non-image files as text
+          fs.writeFileSync(current.name, current.text);
+        }
       } else {
         fs.mkdirSync(current.name, { recursive: true });
       }

package/src/utils/fileDiff.ts

@@ -0,0 +1,114 @@
+import { createTwoFilesPatch } from 'diff';
+import { VizFiles, VizFileId } from '@vizhub/viz-types';
+
+export interface UnifiedFilesDiff {
+  [fileId: VizFileId]: string; // Unified diff string
+}
+
+/**
+ * Generate a unified diff for a single file using the diff library
+ */
+export function generateFileUnifiedDiff(
+  fileId: VizFileId,
+  fileName: string,
+  beforeContent: string,
+  afterContent: string,
+): string {
+  // Only generate diff if there are actual changes
+  if (beforeContent === afterContent) {
+    return '';
+  }
+
+  // Use the diff library's createTwoFilesPatch function to generate unified diff
+  const unifiedDiff = createTwoFilesPatch(
+    fileName,
+    fileName,
+    beforeContent,
+    afterContent,
+  );
+
+  return unifiedDiff;
+}
+
+/**
+ * Generate unified diffs for multiple files
+ */
+export function generateFilesUnifiedDiff(
+  beforeFiles: VizFiles,
+  afterFiles: VizFiles,
+): UnifiedFilesDiff {
+  const result: UnifiedFilesDiff = {};
+  const allFileIds = new Set([
+    ...Object.keys(beforeFiles),
+    ...Object.keys(afterFiles),
+  ]);
+
+  for (const fileId of allFileIds) {
+    const beforeFile = beforeFiles[fileId];
+    const afterFile = afterFiles[fileId];
+
+    const beforeContent = beforeFile?.text || '';
+    const afterContent = afterFile?.text || '';
+    const fileName =
+      afterFile?.name || beforeFile?.name || fileId;
+
+    const unifiedDiff = generateFileUnifiedDiff(
+      fileId,
+      fileName,
+      beforeContent,
+      afterContent,
+    );
+
+    // Only include files that have changes
+    if (unifiedDiff) {
+      result[fileId] = unifiedDiff;
+    }
+  }
+
+  return result;
+}
+
+/**
+ * Create a snapshot of current files for diff comparison
+ */
+export function createFilesSnapshot(
+  files: VizFiles,
+): VizFiles {
+  return JSON.parse(JSON.stringify(files));
+}
+
+/**
+ * Parse unified diff to extract basic statistics
+ */
+export function parseUnifiedDiffStats(
+  unifiedDiff: string,
+): {
+  additions: number;
+  deletions: number;
+} {
+  const lines = unifiedDiff.split('\n');
+  let additions = 0;
+  let deletions = 0;
+
+  for (const line of lines) {
+    if (line.startsWith('+') && !line.startsWith('+++')) {
+      additions++;
+    } else if (
+      line.startsWith('-') &&
+      !line.startsWith('---')
+    ) {
+      deletions++;
+    }
+  }
+
+  return { additions, deletions };
+}
+
+/**
+ * Combine multiple unified diffs into a single diff string
+ */
+export function combineUnifiedDiffs(
+  unifiedDiffs: UnifiedFilesDiff,
+): string {
+  return Object.values(unifiedDiffs).join('\n');
+}