vzcode 1.39.0 → 1.41.0

package/src/server/aiChatHandler/{index.js → index.ts}

@@ -8,12 +8,22 @@ import { createLLMFunction } from './llmStreaming.js';
 import { performAIEditing } from './aiEditing.js';
 import { handleError } from './errorHandling.js';
 import { createRunCodeFunction } from '../../runCode.js';
+import { ShareDBDoc } from '../../types.js';
+import { VizContent } from '@vizhub/viz-types';
 
 const DEBUG = false;
 
 export const handleAIChatMessage =
-  ({ shareDBDoc, localPresence, onCreditDeduction }) =>
-  async (req, res) => {
+  ({
+    shareDBDoc,
+    createVizBotLocalPresence,
+    onCreditDeduction,
+  }: {
+    shareDBDoc: ShareDBDoc<VizContent>;
+    createVizBotLocalPresence: () => any;
+    onCreditDeduction?: any;
+  }) =>
+  async (req: any, res: any) => {
     const { content, chatId } = req.body;
 
     if (DEBUG) {
@@ -41,7 +51,7 @@ export const handleAIChatMessage =
       // Create LLM function for streaming
       const llmFunction = createLLMFunction({
         shareDBDoc,
-        localPresence,
+        createVizBotLocalPresence,
         chatId,
       });
 
@@ -52,7 +62,6 @@ export const handleAIChatMessage =
       const editResult = await performAIEditing({
         prompt: content,
         shareDBDoc,
-        chatId,
         llmFunction,
         runCode,
       });
@@ -60,14 +69,15 @@ export const handleAIChatMessage =
       // Handle credit deduction if callback is provided
       if (
         onCreditDeduction &&
-        editResult.upstreamCostCents
+        (editResult as any).upstreamCostCents
       ) {
         try {
           await onCreditDeduction({
-            upstreamCostCents: editResult.upstreamCostCents,
-            provider: editResult.provider,
-            inputTokens: editResult.inputTokens,
-            outputTokens: editResult.outputTokens,
+            upstreamCostCents: (editResult as any)
+              .upstreamCostCents,
+            provider: (editResult as any).provider,
+            inputTokens: (editResult as any).inputTokens,
+            outputTokens: (editResult as any).outputTokens,
           });
         } catch (creditError) {
           console.error(

package/src/server/aiChatHandler/{llmStreaming.js → llmStreaming.ts}

@@ -16,10 +16,11 @@ const DEBUG = false;
  */
 export const createLLMFunction = ({
   shareDBDoc,
-  localPresence,
+  createVizBotLocalPresence,
   chatId,
 }) => {
   return async (fullPrompt) => {
+    const localPresence = createVizBotLocalPresence();
     // Submit initial presence for VizBot
     const vizBotPresence = {
       username: 'VizBot',
@@ -68,7 +69,10 @@ export const createLLMFunction = ({
 
     // Define callbacks for streaming parser
     const callbacks = {
-      onFileNameChange: (fileName, format) => {
+      onFileNameChange: async (
+        fileName: string,
+        format: string,
+      ) => {
         DEBUG &&
           console.log(
             `File changed to: ${fileName} (${format})`,
@@ -110,7 +114,7 @@ export const createLLMFunction = ({
           'Editing ' + fileName,
         );
       },
-      onCodeLine: (line) => {
+      onCodeLine: async (line: string) => {
         DEBUG && console.log(`Code line: ${line}`);
 
         if (currentEditingFileId) {
@@ -153,7 +157,7 @@ export const createLLMFunction = ({
           }
         }
       },
-      onNonCodeLine: (line) => {
+      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
@@ -175,20 +179,25 @@ export const createLLMFunction = ({
       if (chunk.content) {
         const chunkContent = String(chunk.content);
         chunks.push(chunkContent);
-        parser.processChunk(chunkContent);
+        await parser.processChunk(chunkContent);
       }
 
       if (!generationId && chunk.lc_kwargs?.id) {
         generationId = chunk.lc_kwargs.id;
       }
     }
-    parser.flushRemaining();
+    await parser.flushRemaining();
     reportFileEdited();
     updateAIStatus(shareDBDoc, chatId, 'Done editing.');
     updateAIScratchpad(shareDBDoc, chatId, fullContent);
 
     // Clear VizBot presence when done
-    localPresence.destroy((error) => {
+    DEBUG &&
+      console.log(
+        'AI editing done, clearing VizBot presence',
+      );
+    localPresence.submit(null, (error) => {
+      DEBUG && console.log('VizBot presence cleared');
       if (error) {
         console.warn(
           'VizBot presence cleanup error:',

package/src/server/{generateAIResponse.js → generateAIResponse.ts}

@@ -28,7 +28,8 @@ const slowdown = false;
 
 // The options passed into the OpenAI client
 // new OpenAI(openAIOptions)
-const openAIOptions = {};
+const openAIOptions: { apiKey?: string; baseURL?: string } =
+  {};
 
 // Support specifying the API key via an environment variable
 // If VZCODE_AI_API_KEY is not set, note that the OpenAI client
@@ -71,7 +72,7 @@ const opComesFromAIAssist = (ops, source) =>
 
 // Keeps track of the currently ongoing AI streams.
 // There could be many streams at the same time.
-const streams = {};
+export const streams: Record<string, any> = {};
 
 export const generateAIResponse = async ({
   inputText,

package/src/server/{handleAIAssist.js → handleAIAssist.ts}

@@ -1,9 +1,12 @@
-import { generateAIResponse } from './generateAIResponse.js';
+import {
+  generateAIResponse,
+  streams,
+} from './generateAIResponse.js';
 
 const debug = false;
 
 export const handleAIAssist =
-  (shareDBDoc) => async (req, res) => {
+  (shareDBDoc: any) => async (req: any, res: any) => {
     const { inputText, insertionCursor, fileId } = req.body;
 
     if (debug) {
@@ -16,17 +19,21 @@ export const handleAIAssist =
     }
 
     try {
+      // Generate a unique streamId for this request
+      const streamId = `stream_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
+
       await generateAIResponse({
         inputText,
         insertionCursor,
         fileId,
+        streamId,
         shareDBDoc,
       });
 
       res
         .status(200)
         .send({ message: 'Operation successful!' });
-    } catch (error) {
+    } catch (error: any) {
       console.error('handleAIAssist error:', error);
       res.status(500).send({
         message: 'Internal Server Error',
@@ -35,7 +42,7 @@ export const handleAIAssist =
     }
   };
 
-export function haltGeneration(streamId) {
+export function haltGeneration(streamId: string) {
   const stream = streams[streamId];
 
   // Stream can be undefined here if the user

package/src/server/{index.js → index.ts}

@@ -4,7 +4,6 @@ import bodyParser from 'body-parser';
 import express from 'express';
 import fs from 'fs';
 import http from 'http';
-import ngrok from 'ngrok';
 import open from 'open';
 import path from 'path';
 import ShareDB from 'sharedb';
@@ -103,9 +102,8 @@ const generateVizBotId = () => {
   return `vizbot-${timestamp}`;
 };
 
-const localPresence = docPresence.create(
-  generateVizBotId(),
-);
+const createVizBotLocalPresence = () =>
+  docPresence.create(generateVizBotId());
 
 shareDBDoc.create(initialDocument, json1Presence.type.uri);
 
@@ -127,7 +125,11 @@ app.post(
 app.post(
   '/ai-chat-message',
   bodyParser.json(),
-  handleAIChatMessage({ shareDBDoc, localPresence }),
+  handleAIChatMessage({
+    shareDBDoc,
+    createVizBotLocalPresence,
+    onCreditDeduction: undefined,
+  }),
 );
 
 // Livekit Token Generator
@@ -335,21 +337,13 @@ shareDBDoc.subscribe(() => {
 });
 
 server.listen(port, async () => {
-  if (process.env.NGROK_TOKEN) {
-    (async function () {
-      await ngrok.authtoken(process.env.NGROK_TOKEN);
-      const url = await ngrok.connect(port);
-      console.log(`Editor is live at ${url}`);
-      open(url);
-    })();
-  } else {
-    // Sets the port to the one specified in the environment
-    // variable (for development) or the default port.
-    let livePort = process.env.EDITOR_PORT || port;
-    console.log(`EDITOR_PORT: ${process.env.EDITOR_PORT}`);
-    console.log(
-      `Editor is live at http://localhost:${livePort}`,
-    );
-    open(`http://localhost:${livePort}`);
-  }
+  // Note: ngrok support can be added when the package is properly installed
+  // Sets the port to the one specified in the environment
+  // variable (for development) or the default port.
+  let livePort = process.env.EDITOR_PORT || port;
+  console.log(`EDITOR_PORT: ${process.env.EDITOR_PORT}`);
+  console.log(
+    `Editor is live at http://localhost:${livePort}`,
+  );
+  open(`http://localhost:${livePort}`);
 });