vzcode 2.3.0 → 2.4.0

package/src/server/aiChatHandler/index.ts

@@ -0,0 +1,199 @@
+import { validateRequest } from './validation.js';
+import {
+  ensureChatsExist,
+  ensureChatExists,
+  addUserMessage,
+  addDiffToAIMessage,
+  setAIStatus,
+} from './chatOperations.js';
+import { createLLMFunction } from './llmStreaming.js';
+import {
+  performAIEditing,
+  performAIChat,
+} from './aiEditing.js';
+import {
+  handleError,
+  handleBackgroundError,
+} from './errorHandling.js';
+import { createRunCodeFunction } from '../../runCode.js';
+import { ShareDBDoc } from '../../types.js';
+import { VizContent } from '@vizhub/viz-types';
+import { createSubmitOperation } from '../../submitOperation.js';
+
+const DEBUG = false;
+
+export const handleAIChatMessage =
+  ({
+    shareDBDoc,
+    createAIEditLocalPresence,
+    onCreditDeduction,
+    getCurrentCommitId,
+    model,
+    aiRequestOptions,
+  }: {
+    shareDBDoc: ShareDBDoc<VizContent>;
+    createAIEditLocalPresence: () => any;
+    onCreditDeduction?: any;
+    getCurrentCommitId?: () => string | null;
+    model?: string;
+    aiRequestOptions?: any;
+  }) =>
+  async (req: any, res: any) => {
+    const { content, chatId, mode = 'edit' } = req.body;
+
+    if (DEBUG) {
+      console.log(
+        '[handleAIChatMessage] content:',
+        content,
+        'chatId:',
+        chatId,
+        'shareDBDoc:',
+        shareDBDoc,
+      );
+    }
+
+    // Validate request
+    if (!validateRequest(req, res)) {
+      return;
+    }
+
+    try {
+      // Ensure chats structure exists
+      ensureChatsExist(shareDBDoc);
+      ensureChatExists(shareDBDoc, chatId);
+
+      // Add user message to chat
+      addUserMessage(shareDBDoc, chatId, content);
+
+      // Return success immediately - AI generation continues in background
+      res.status(200).json('success');
+
+      // Set AI status to indicate generation is starting
+      setAIStatus(shareDBDoc, chatId, 'generating');
+
+      // Continue AI processing in background (don't await)
+      processAIRequestAsync({
+        shareDBDoc,
+        chatId,
+        content,
+        mode,
+        createAIEditLocalPresence,
+        getCurrentCommitId,
+        model,
+        aiRequestOptions,
+        onCreditDeduction,
+      }).catch((error) => {
+        console.error(
+          'Background AI processing error:',
+          error,
+        );
+        // Handle error without HTTP response
+        handleBackgroundError(shareDBDoc, chatId, error);
+      });
+    } catch (error) {
+      handleError(shareDBDoc, chatId, error, res);
+    }
+  };
+
+/**
+ * Processes the AI request asynchronously in the background
+ */
+const processAIRequestAsync = async ({
+  shareDBDoc,
+  chatId,
+  content,
+  mode,
+  createAIEditLocalPresence,
+  getCurrentCommitId,
+  model,
+  aiRequestOptions,
+  onCreditDeduction,
+}: {
+  shareDBDoc: ShareDBDoc<VizContent>;
+  chatId: string;
+  content: string;
+  mode: string;
+  createAIEditLocalPresence: () => any;
+  getCurrentCommitId?: () => string | null;
+  model?: string;
+  aiRequestOptions?: any;
+  onCreditDeduction?: any;
+}) => {
+  try {
+    // Capture the current commit ID before making changes (for VizHub integration)
+    const beforeCommitId = getCurrentCommitId
+      ? getCurrentCommitId()
+      : null;
+
+    // Create LLM function for streaming
+    const llmFunction = createLLMFunction({
+      shareDBDoc,
+      createAIEditLocalPresence,
+      chatId,
+      model,
+      aiRequestOptions,
+    });
+
+    // Create server-side runCode function using shareDBDoc
+    const submitOperation =
+      createSubmitOperation(shareDBDoc);
+    const runCode = createRunCodeFunction(submitOperation);
+
+    // 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 (
+      editResult.diffData &&
+      Object.keys(editResult.diffData).length > 0
+    ) {
+      addDiffToAIMessage(
+        shareDBDoc,
+        chatId,
+        editResult.diffData,
+        beforeCommitId, // Pass the commit ID before AI changes (for VizHub integration)
+      );
+    }
+
+    // Handle credit deduction if callback is provided
+    if (
+      onCreditDeduction &&
+      (editResult as any).upstreamCostCents
+    ) {
+      try {
+        await onCreditDeduction({
+          upstreamCostCents: (editResult as any)
+            .upstreamCostCents,
+          provider: (editResult as any).provider,
+          inputTokens: (editResult as any).inputTokens,
+          outputTokens: (editResult as any).outputTokens,
+        });
+      } catch (creditError) {
+        console.error(
+          'Credit deduction error:',
+          creditError,
+        );
+        // Don't fail the request if credit deduction fails
+      }
+    }
+
+    // Clear the AI status to indicate completion
+    setAIStatus(shareDBDoc, chatId, undefined);
+  } catch (error) {
+    // Set error status and add error message to chat
+    setAIStatus(shareDBDoc, chatId, 'error');
+    handleBackgroundError(shareDBDoc, chatId, error);
+  }
+};

package/src/server/aiChatHandler/llmStreaming.ts

@@ -0,0 +1,419 @@
+import {
+  parseMarkdownFiles,
+  StreamingMarkdownParser,
+} from 'llm-code-format';
+import OpenAI from 'openai';
+import fs from 'fs';
+import { generateRunId } from '@vizhub/viz-utils';
+import {
+  updateAIStatus,
+  createAIMessage,
+  updateAIMessageContent,
+  finalizeAIMessage,
+  ensureFileExists,
+  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;
+
+// Useful for testing/debugging the streaming behavior
+const slowMode = false;
+
+// Throttle the streaming updates, so that we don't
+// overwhelm the ShareDB server with too many updates.
+// It happened actually, before adding this.
+// MongoDB VizHub server got in fact overloaded with
+// too many updates from the AI streaming response, with
+// warning: "Replication Oplog Window has gone below 1 hour"
+const THROTTLE_INTERVAL_MS = 100;
+
+// Feature flag to enable/disable streaming editing.
+// * If `true`, the AI streaming response will be used to
+//   edit files in real-time by submitting ShareDB ops.
+// * If `false`, the updates to code files will be applied
+// only after the AI has finished generating the entire response.
+//
+// Current status: there's a tricky bug with the streaming
+// where the AI edits sometimes don't apply correctly in CodeMirror.
+// It seems that sometimes the op that clears the file content
+// is not applied correctly in the front end, leading to
+// a situation where the AI streaming edits are concatenated into the middle
+// of the file instead of replacing it.
+// See https://github.com/codemirror/codemirror.next/issues/1234
+const enableStreamingEditing = false;
+
+/**
+ * Creates and configures the LLM function for streaming with reasoning tokens
+ */
+export const createLLMFunction = ({
+  shareDBDoc,
+  createAIEditLocalPresence,
+  chatId,
+  // Feature flag to enable/disable reasoning tokens.
+  // When false, reasoning tokens are not requested from the API
+  // and reasoning content is not processed in the streaming response.
+  enableReasoningTokens = false,
+  model,
+  aiRequestOptions,
+}: {
+  shareDBDoc: ShareDBDoc<VizContent>;
+  createAIEditLocalPresence: () => any;
+  chatId: VizChatId;
+  enableReasoningTokens?: boolean;
+  model?: string;
+  aiRequestOptions?: any;
+}) => {
+  return async (fullPrompt: string) => {
+    const localPresence = enableStreamingEditing
+      ? createAIEditLocalPresence()
+      : null;
+
+    // Create OpenRouter client for reasoning token support
+    const openRouterClient = new OpenAI({
+      apiKey: process.env.VZCODE_EDIT_WITH_AI_API_KEY,
+      baseURL:
+        process.env.VZCODE_EDIT_WITH_AI_BASE_URL ||
+        'https://openrouter.ai/api/v1',
+      defaultHeaders: {
+        'HTTP-Referer': 'https://vizhub.com',
+        'X-Title': 'VizHub',
+      },
+    });
+
+    let fullContent = '';
+    let generationId = '';
+    let currentEditingFileId = null;
+    let currentEditingFileName = null;
+
+    // Create initial AI message for streaming
+    const aiMessageId = createAIMessage(shareDBDoc, chatId);
+
+    // --- throttle wrapper ---
+    function makeThrottledUpdater() {
+      let lastCall = 0;
+      let latestContent = '';
+      let timer: NodeJS.Timeout | null = null;
+
+      function invoke() {
+        updateAIMessageContent(
+          shareDBDoc,
+          chatId,
+          aiMessageId,
+          latestContent,
+        );
+        lastCall = Date.now();
+      }
+
+      const fn = (content: string) => {
+        latestContent = content;
+        const now = Date.now();
+
+        if (now - lastCall >= THROTTLE_INTERVAL_MS) {
+          // safe to call immediately
+          invoke();
+        } else if (!timer) {
+          // schedule for later
+          timer = setTimeout(
+            () => {
+              timer = null;
+              invoke();
+            },
+            THROTTLE_INTERVAL_MS - (now - lastCall),
+          );
+        }
+      };
+
+      // expose a flush() helper to force an immediate write
+      fn.flush = () => {
+        if (timer) {
+          clearTimeout(timer);
+          timer = null;
+        }
+        invoke();
+      };
+
+      return fn;
+    }
+
+    const throttledUpdateAIMessageContent =
+      makeThrottledUpdater();
+
+    // Function to report file edited
+    // This is called when the AI has finished editing a file
+    // and we want to update the message content with the file name.
+    const reportFileEdited = () => {
+      if (currentEditingFileName) {
+        fullContent += ` * Edited ${currentEditingFileName}\n`;
+        throttledUpdateAIMessageContent(fullContent);
+        currentEditingFileName = null;
+      }
+    };
+
+    // Define callbacks for streaming parser
+    const callbacks = {
+      onFileNameChange: async (
+        fileName: string,
+        format: string,
+      ) => {
+        DEBUG &&
+          console.log(
+            `File changed to: ${fileName} (${format})`,
+          );
+
+        // Find existing file or create new one
+        currentEditingFileId = ensureFileExists(
+          shareDBDoc,
+          fileName,
+        );
+
+        reportFileEdited();
+        currentEditingFileName = fileName;
+
+        // Clear the file content to start fresh
+        // (AI will regenerate the entire file content)
+        clearFileContent(shareDBDoc, currentEditingFileId);
+
+        // Update AI status
+        updateAIStatus(
+          shareDBDoc,
+          chatId,
+          'Editing ' + fileName,
+        );
+      },
+      onCodeLine: async (line: string) => {
+        DEBUG && console.log(`Code line: ${line}`);
+
+        // If streaming is enabled, we apply the line immediately
+        if (currentEditingFileId) {
+          // Apply OT operation for this line immediately
+          appendLineToFile(
+            shareDBDoc,
+            currentEditingFileId,
+            line,
+          );
+
+          // Update AI presence to show cursor at the end of the file
+          const currentFile =
+            shareDBDoc.data.files[currentEditingFileId];
+          if (currentFile && currentFile.text) {
+            const textLength = currentFile.text.length;
+            const filePresence = {
+              username: 'AI Editor',
+              start: [
+                'files',
+                currentEditingFileId,
+                'text',
+                textLength,
+              ],
+              end: [
+                'files',
+                currentEditingFileId,
+                'text',
+                textLength,
+              ],
+            };
+
+            if (localPresence) {
+              localPresence.submit(
+                filePresence,
+                (error) => {
+                  if (error) {
+                    console.warn(
+                      'AI Editor line presence submission error:',
+                      error,
+                    );
+                  }
+                },
+              );
+            }
+          }
+        }
+      },
+      onNonCodeLine: async (line: string) => {
+        // We want to report a file edited only if the line is not empty,
+        // because sometimes the LLMs leave a newline between the file name
+        // declaration and th
+        if (line.trim() !== '') {
+          reportFileEdited();
+        }
+        fullContent += line + '\n';
+        throttledUpdateAIMessageContent(fullContent);
+      },
+    };
+
+    const parser = new StreamingMarkdownParser(callbacks);
+
+    const chunks = [];
+    let reasoningContent = '';
+
+    // Stream the response with reasoning tokens
+    const modelName =
+      model ||
+      process.env.VZCODE_EDIT_WITH_AI_MODEL_NAME ||
+      'anthropic/claude-3.5-sonnet';
+
+    // Configure reasoning tokens based on enableReasoningTokens flag
+    const requestConfig: any = {
+      model: modelName,
+      messages: [{ role: 'user', content: fullPrompt }],
+      usage: { include: true },
+      stream: true,
+      ...aiRequestOptions,
+    };
+
+    // Only include reasoning configuration if reasoning tokens are enabled
+    if (enableReasoningTokens) {
+      requestConfig.reasoning = {
+        effort: 'low',
+        exclude: false,
+      };
+    }
+
+    const stream = await (
+      openRouterClient.chat.completions.create as any
+    )(requestConfig);
+
+    let reasoningStarted = false;
+    let contentStarted = false;
+
+    for await (const chunk of stream) {
+      if (slowMode) {
+        await new Promise((resolve) =>
+          setTimeout(resolve, 500),
+        );
+      }
+      const delta = chunk.choices[0]?.delta as any; // Type assertion for OpenRouter-specific reasoning fields
+
+      if (delta?.reasoning && enableReasoningTokens) {
+        // Handle reasoning tokens (thinking) - only if enabled
+        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) {
+          await parser.processChunk(chunkContent);
+        } else {
+          fullContent += chunkContent;
+          throttledUpdateAIMessageContent(fullContent);
+        }
+      } else if (chunk.usage) {
+        // Handle usage information
+        DEBUG && console.log('Usage:', chunk.usage);
+      }
+
+      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.');
+    throttledUpdateAIMessageContent(fullContent);
+
+    // Flush to ensure the final content is written immediately
+    throttledUpdateAIMessageContent.flush();
+
+    // Finalize the AI message by clearing temporary fields
+    finalizeAIMessage(shareDBDoc, chatId);
+
+    // Clear AI Editor presence when done
+    DEBUG &&
+      console.log(
+        'AI editing done, clearing AI Editor presence',
+      );
+    if (localPresence) {
+      localPresence.submit(null, (error) => {
+        DEBUG && console.log('AI Editor presence cleared');
+        if (error) {
+          console.warn(
+            'AI Editor presence cleanup error:',
+            error,
+          );
+        }
+      });
+    }
+
+    // If streaming editing is not enabled, we need to
+    // apply all the edits at once
+    if (!enableStreamingEditing) {
+      updateFiles(
+        shareDBDoc,
+        mergeFileChanges(
+          shareDBDoc.data.files,
+          parseMarkdownFiles(fullContent, 'bold').files,
+        ),
+      );
+    }
+
+    // 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`;
+      fs.writeFileSync(
+        chunksFileJSONpath,
+        JSON.stringify(chunks, null, 2),
+      );
+      console.log(
+        `AI chunks written to ${chunksFileJSONpath}`,
+      );
+    }
+
+    return {
+      content: fullContent,
+      generationId: generationId,
+    };
+  };
+};

package/src/server/aiChatHandler/validation.ts

@@ -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;
+};