vzcode 1.45.0 → 1.47.0

package/src/server/aiChatHandler/chatOperations.ts

@@ -1,4 +1,7 @@
-import { dateToTimestamp } from '@vizhub/viz-utils';
+import {
+  dateToTimestamp,
+  generateRunId,
+} from '@vizhub/viz-utils';
 import { randomId } from '../../randomId.js';
 import { ShareDBDoc } from '../../types.js';
 import { diff } from '../../ot.js';
@@ -253,6 +256,7 @@ export const addDiffToAIMessage = (
   shareDBDoc: ShareDBDoc<VizContent>,
   chatId: VizChatId,
   diffData: any,
+  beforeFiles?: VizFiles, // Optional snapshot for undo functionality
 ) => {
   const chat = shareDBDoc.data.chats[chatId];
   const messages = [...chat.messages];
@@ -266,6 +270,7 @@ export const addDiffToAIMessage = (
     const newMessage = {
       ...messages[lastAIMessageIndex],
       diffData,
+      ...(beforeFiles && { beforeFiles }), // Add beforeFiles only if provided
     };
     // Use type assertion to extend the message with diffData
     (messages[lastAIMessageIndex] as any) = newMessage;
@@ -447,3 +452,35 @@ export const appendLineToFile = (
 
   shareDBDoc.submitOp(diff(shareDBDoc.data, newDocState));
 };
+
+/**
+ * Undoes the last AI edit by restoring files and removing the AI message
+ */
+export const undoAIEdit = (
+  shareDBDoc: ShareDBDoc<VizContent>,
+  chatId: VizChatId,
+  messageId: string,
+  beforeFiles: VizFiles,
+) => {
+  const chat = shareDBDoc.data.chats[chatId];
+  const messages = chat.messages.filter(
+    (msg) => msg.id !== messageId,
+  );
+
+  const op = diff(shareDBDoc.data, {
+    ...shareDBDoc.data,
+    files: beforeFiles,
+    chats: {
+      ...shareDBDoc.data.chats,
+      [chatId]: {
+        ...chat,
+        messages,
+        updatedAt: dateToTimestamp(new Date()),
+      },
+    },
+    // Trigger a re-run
+    runId: generateRunId(),
+  });
+
+  shareDBDoc.submitOp(op);
+};

package/src/server/aiChatHandler/index.ts

@@ -6,7 +6,10 @@ import {
   addDiffToAIMessage,
 } from './chatOperations.js';
 import { createLLMFunction } from './llmStreaming.js';
-import { performAIEditing } from './aiEditing.js';
+import {
+  performAIEditing,
+  performAIChat,
+} from './aiEditing.js';
 import { handleError } from './errorHandling.js';
 import { createRunCodeFunction } from '../../runCode.js';
 import { ShareDBDoc } from '../../types.js';
@@ -26,7 +29,7 @@ export const handleAIChatMessage =
     onCreditDeduction?: any;
   }) =>
   async (req: any, res: any) => {
-    const { content, chatId } = req.body;
+    const { content, chatId, mode = 'edit' } = req.body;
 
     if (DEBUG) {
       console.log(
@@ -65,13 +68,20 @@ export const handleAIChatMessage =
       const runCode =
         createRunCodeFunction(submitOperation);
 
-      // Perform AI editing
-      const editResult = await performAIEditing({
-        prompt: content,
-        shareDBDoc,
-        llmFunction,
-        runCode,
-      });
+      // 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 (
@@ -82,6 +92,7 @@ export const handleAIChatMessage =
           shareDBDoc,
           chatId,
           editResult.diffData,
+          (editResult as any).beforeFiles, // Pass the beforeFiles snapshot for undo (only available for edit mode)
         );
       }
 

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/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,6 +14,7 @@ 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';
@@ -141,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;