Merge pull request #624 from encorelab/623-bug-ai-assistant-error-par…

…sing-json-special-characters

Added server-side escaping of AI Assistant generated JSON special characters

backend/src/services/vertexAI/index.ts

@@ -65,7 +65,8 @@ function checkKeyFile(keyfilePath: string): boolean {
 }
 
 // Check the keyfile before initializing the client
-const keyfilePath = process.env.GOOGLE_APPLICATION_CREDENTIALS || './secrets/keyfile.json';
+const keyfilePath =
+  process.env.GOOGLE_APPLICATION_CREDENTIALS || './secrets/keyfile.json';
 checkKeyFile(keyfilePath);
 
 // Function to initialize the EndpointServiceClient with delayed checking
@@ -181,9 +182,8 @@ function parseVertexAIError(errorString: string): ErrorInfo {
   }
 }
 
-function isValidJSON(str: string): boolean {
+function isValidJSON(jsonObject: any): boolean {
   try {
-    const jsonObject = JSON.parse(str);
 
     if (!jsonObject.response) {
       return false;
@@ -284,8 +284,61 @@ function postsToKeyValuePairs(posts: any[]): string {
 }
 
 function removeJsonMarkdown(text: string): string {
-  const pattern = /```json\s*([\s\S]*?)\s*```/g;
-  return text.replace(pattern, '$1');
+  text = text.trim();
+  const startIndex = text.indexOf('```json');
+  const endIndex = text.lastIndexOf('```');
+
+  if (startIndex === -1 || endIndex === -1) {
+    console.warn('Invalid JSON markdown format:', text);
+    return text; // Or handle the error differently
+  }
+
+  return text.slice(startIndex + '```json'.length, endIndex);
+}
+
+function parseJsonResponse(response: string): any {
+  if (!response) {
+    return {};
+  }
+
+  // Remove the "response" key and its value, including "<END>", "<END>"", or "<END>","
+  const responseStartIndex = response.indexOf('"response": "');
+  let endLength = '<END>"'.length;
+  let responseEndIndex = response.indexOf('<END>",'); 
+
+  if (responseEndIndex === -1) {
+    responseEndIndex = response.indexOf('<END>"');
+    if (responseEndIndex !== -1) {
+      endLength = '<END>"'.length
+      responseEndIndex += endLength; 
+    }
+  } else {
+    endLength = '<END>",'.length
+    responseEndIndex += endLength;
+  }
+
+  if (responseStartIndex === -1 || responseEndIndex === -1) {
+    console.warn('Invalid response format:', response);
+    return {}; 
+  }
+
+  const responseValue = response.substring(
+    responseStartIndex + '"response": "'.length,
+    responseEndIndex + ('<END>'.length - endLength)
+  );
+
+  const textWithoutResponse =
+    response.substring(0, responseStartIndex) +
+    response.substring(responseEndIndex);
+
+  try {
+    const jsonObject = JSON.parse(textWithoutResponse);
+    jsonObject.response = responseValue; // Add back with <END> included
+    return jsonObject;
+  } catch (error) {
+    console.error('Error parsing JSON:', error);
+    return {}; 
+  }
 }
 
 async function sendMessage(
@@ -337,7 +390,6 @@ async function sendMessage(
           // Async IIFE
           for await (const item of stream) {
             partialResponse += item.candidates[0].content.parts[0].text;
-            // console.log("Partial response:", partialResponse);
 
             socket.emit(SocketEvent.AI_RESPONSE, {
               status: 'Processing',
@@ -346,11 +398,12 @@ async function sendMessage(
           }
 
           let isValid;
+          const noJsonResponse = removeJsonMarkdown(partialResponse);
+
           try {
-            const cleanedResponse = removeJsonMarkdown(partialResponse);
-            const parsedResponse = JSON.parse(cleanedResponse);
+            const parsedResponse = parseJsonResponse(noJsonResponse)
 
-            if (isValidJSON(cleanedResponse)) {
+            if (isValidJSON(parsedResponse)) {
               finalResponse = parsedResponse;
               isValid = true;
             } else {
@@ -367,9 +420,7 @@ async function sendMessage(
               response: finalResponse.response,
             });
           } else {
-            const errorMessage =
-              'Invalid response formatting. Please try again.\n\n' +
-              finalResponse.response;
+            const errorMessage = `Completed with invalid formatting: ${partialResponse}`;
             socket.emit(SocketEvent.AI_RESPONSE, {
               status: 'Error',
               errorMessage: errorMessage,

frontend/src/app/components/create-workflow-modal/create-workflow-modal.component.ts

@@ -450,8 +450,8 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
                 break;
               }
               case 'Completed': {
-                const escapedResponse = this.escapeJsonResponse(data.response);
-                this.aiResponse = this.markdownToHtml(escapedResponse || '');
+                const dataResponse = data.response;
+                this.aiResponse = this.markdownToHtml(dataResponse || '');
                 this.chatHistory.push({
                   role: 'assistant',
                   content: this.aiResponse,
@@ -499,33 +499,6 @@ export class CreateWorkflowModalComponent implements OnInit, OnDestroy {
     this.isProcessingAIRequest = false;
   }
 
-  private escapeJsonResponse(response: string): string {
-    if (!response) {
-      return '';
-    }
-
-    const responseStartIndex =
-      response.indexOf('"response": "') + '"response": "'.length;
-    const responseEndIndex = response.indexOf('<END>');
-
-    if (responseStartIndex === -1 || responseEndIndex === -1) {
-      console.warn('Invalid response format:', response);
-      return response; // Or handle the error differently
-    }
-
-    const responseValue = response.substring(
-      responseStartIndex,
-      responseEndIndex
-    );
-    const escapedValue = JSON.stringify(responseValue).slice(1, -1); // Escape special characters
-
-    return (
-      response.substring(0, responseStartIndex) +
-      escapedValue +
-      response.substring(responseEndIndex)
-    );
-  }
-
   // Method to scroll the div to the bottom
   scrollToBottom(): void {
     const scrollableElement = this.scrollableDiv.nativeElement;