@probelabs/probe 0.6.0-rc163 → 0.6.0-rc165

package/build/agent/schemaUtils.js

@@ -63,12 +63,6 @@ export function generateSchemaInstructions(schema, options = {}) {
   try {
     const parsedSchema = typeof schema === 'string' ? JSON.parse(schema) : schema;
     instructions += `${JSON.stringify(parsedSchema, null, 2)}\n\n`;
-
-    // Generate example using helper function
-    const exampleObj = generateExampleFromSchema(parsedSchema, { debug });
-    if (exampleObj) {
-      instructions += `Example:\n<attempt_completion>\n${JSON.stringify(exampleObj, null, 2)}\n</attempt_completion>\n\n`;
-    }
   } catch (e) {
     if (debug) {
       console.error('[DEBUG] generateSchemaInstructions: Failed to parse schema:', e.message);
@@ -76,7 +70,7 @@ export function generateSchemaInstructions(schema, options = {}) {
     instructions += `${schema}\n\n`;
   }
 
-  instructions += 'Your response inside attempt_completion must be ONLY valid JSON - no plain text, no explanations, no markdown.';
+  instructions += 'Your response inside attempt_completion must be ONLY valid JSON - no plain text, no explanations, no markdown.\n\nIMPORTANT: First complete the requested analysis/task thoroughly, then provide your final answer in the JSON format above.';
 
   return instructions;
 }
@@ -171,6 +165,74 @@ export function decodeHtmlEntities(text) {
   return decoded;
 }
 
+/**
+ * Normalize JavaScript syntax to valid JSON syntax
+ * Converts single quotes to double quotes for strings in JSON-like structures
+ *
+ * @param {string} str - String that might contain JavaScript array/object syntax
+ * @returns {string} - String with single quotes normalized to double quotes
+ */
+function normalizeJsonQuotes(str) {
+  if (!str || typeof str !== 'string') {
+    return str;
+  }
+
+  // Quick check: if there are no single quotes, no need to normalize
+  if (!str.includes("'")) {
+    return str;
+  }
+
+  let result = '';
+  let inDoubleQuote = false;
+  let inSingleQuote = false;
+  let escaped = false;
+
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i];
+    const prevChar = i > 0 ? str[i - 1] : '';
+
+    // Handle escape sequences
+    if (escaped) {
+      result += char;
+      escaped = false;
+      continue;
+    }
+
+    if (char === '\\') {
+      escaped = true;
+      result += char;
+      continue;
+    }
+
+    // Track when we're inside double-quoted strings
+    if (char === '"' && !inSingleQuote) {
+      inDoubleQuote = !inDoubleQuote;
+      result += char;
+      continue;
+    }
+
+    // Convert single quotes to double quotes (when not inside double quotes)
+    if (char === "'" && !inDoubleQuote) {
+      // Check if this is a single quote inside a string value (like "It's")
+      // If we're already in a single-quoted string, toggle the state
+      if (inSingleQuote) {
+        // Closing single quote - convert to double quote
+        result += '"';
+        inSingleQuote = false;
+      } else {
+        // Opening single quote - convert to double quote
+        result += '"';
+        inSingleQuote = true;
+      }
+      continue;
+    }
+
+    result += char;
+  }
+
+  return result;
+}
+
 /**
  * Clean AI response by extracting JSON content when response contains JSON
  * Only processes responses that contain JSON structures { or [
@@ -195,29 +257,36 @@ export function cleanSchemaResponse(response) {
   // Try with json language specifier
   const jsonBlockMatch = trimmed.match(/```json\s*\n([\s\S]*?)\n```/);
   if (jsonBlockMatch) {
-    return jsonBlockMatch[1].trim();
+    return normalizeJsonQuotes(jsonBlockMatch[1].trim());
   }
 
   // Try any code block with JSON content
   const anyBlockMatch = trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
   if (anyBlockMatch) {
-    return anyBlockMatch[1].trim();
+    return normalizeJsonQuotes(anyBlockMatch[1].trim());
   }
 
   // Legacy patterns for more specific matching
   const codeBlockPatterns = [
     /```json\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /`([{\[][\s\S]*?[}\]])`/
+    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/
   ];
 
   for (const pattern of codeBlockPatterns) {
     const match = trimmed.match(pattern);
     if (match) {
-      return match[1].trim();
+      return normalizeJsonQuotes(match[1].trim());
     }
   }
 
+  // Single backtick pattern - ONLY if the entire input is just the code block
+  // This prevents extracting inline code from within JSON objects (e.g., `['*']` from markdown text)
+  const singleBacktickPattern = /^`([{\[][\s\S]*?[}\]])`$/;
+  const singleBacktickMatch = trimmed.match(singleBacktickPattern);
+  if (singleBacktickMatch) {
+    return normalizeJsonQuotes(singleBacktickMatch[1].trim());
+  }
+
   // Look for code block start followed immediately by JSON
   const codeBlockStartPattern = /```(?:json)?\s*\n?\s*([{\[])/;
   const codeBlockMatch = trimmed.match(codeBlockStartPattern);
@@ -242,7 +311,7 @@ export function cleanSchemaResponse(response) {
     }
 
     if (bracketCount === 0) {
-      return trimmed.substring(startIndex, endIndex);
+      return normalizeJsonQuotes(trimmed.substring(startIndex, endIndex));
     }
   }
 
@@ -267,7 +336,8 @@ export function cleanSchemaResponse(response) {
   const isJsonArray = firstChar === '[' && lastChar === ']';
 
   if (isJsonObject || isJsonArray) {
-    return cleaned;
+    // Normalize JavaScript syntax (single quotes) to valid JSON syntax (double quotes)
+    return normalizeJsonQuotes(cleaned);
   }
 
   return response; // Return original if no extractable JSON found
@@ -832,6 +902,76 @@ export function isMermaidSchema(schema) {
   return mermaidIndicators.some(indicator => indicator);
 }
 
+/**
+ * Extract Mermaid diagrams from JSON string values
+ * Handles escaped newlines and backticks within JSON strings
+ * @param {string} response - Response that may contain JSON with mermaid blocks in string values
+ * @returns {Object} - {diagrams: Array, jsonPaths: Array, parsedJson: Object|null}
+ */
+export function extractMermaidFromJson(response) {
+  if (!response || typeof response !== 'string') {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+
+  // Try to extract JSON from code blocks first
+  let jsonContent = response.trim();
+  const jsonBlockMatch = jsonContent.match(/```json\s*\n([\s\S]*?)\n```/);
+  if (jsonBlockMatch) {
+    jsonContent = jsonBlockMatch[1].trim();
+  } else {
+    const anyBlockMatch = jsonContent.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
+    if (anyBlockMatch) {
+      jsonContent = anyBlockMatch[1].trim();
+    }
+  }
+
+  // Try to parse as JSON
+  let parsedJson;
+  try {
+    parsedJson = JSON.parse(jsonContent);
+  } catch (e) {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+
+  const diagrams = [];
+  const jsonPaths = [];
+
+  // Recursively search for mermaid diagrams in JSON string values
+  function searchObject(obj, path = []) {
+    if (typeof obj === 'string') {
+      // Look for mermaid code blocks in the string value
+      // Handle both escaped (\n) and literal newlines
+      const mermaidPattern = /```mermaid([^\n`]*?)(?:\n|\\n)([\s\S]*?)```/gi;
+      let match;
+
+      while ((match = mermaidPattern.exec(obj)) !== null) {
+        const attributes = match[1] ? match[1].trim() : '';
+        // Unescape the content (replace \\n with actual newlines)
+        const content = match[2].replace(/\\n/g, '\n');
+
+        diagrams.push({
+          content: content,
+          fullMatch: match[0],
+          startIndex: match.index,
+          endIndex: match.index + match[0].length,
+          attributes: attributes,
+          isInJson: true,
+          jsonPath: path.join('.')
+        });
+        jsonPaths.push(path.join('.'));
+      }
+    } else if (Array.isArray(obj)) {
+      obj.forEach((item, index) => searchObject(item, [...path, `[${index}]`]));
+    } else if (obj && typeof obj === 'object') {
+      Object.entries(obj).forEach(([key, value]) => searchObject(value, [...path, key]));
+    }
+  }
+
+  searchObject(parsedJson);
+
+  return { diagrams, jsonPaths, parsedJson };
+}
+
 /**
  * Extract Mermaid diagrams from markdown code blocks with position tracking
  * @param {string} response - Response that may contain markdown with mermaid blocks
@@ -842,6 +982,16 @@ export function extractMermaidFromMarkdown(response) {
     return { diagrams: [], cleanedResponse: response };
   }
 
+  // First check if this looks like a JSON response - if so, use JSON-aware extraction
+  const trimmed = response.trim();
+  if ((trimmed.startsWith('{') || trimmed.startsWith('[')) ||
+      trimmed.includes('```json')) {
+    const jsonResult = extractMermaidFromJson(response);
+    if (jsonResult.diagrams.length > 0) {
+      return { diagrams: jsonResult.diagrams, cleanedResponse: response };
+    }
+  }
+
   // Find all mermaid code blocks with enhanced regex to capture more variations
   // This regex captures optional attributes on same line as ```mermaid, and all diagram content
   const mermaidBlockRegex = /```mermaid([^\n]*)\n([\s\S]*?)```/gi;
@@ -860,7 +1010,8 @@ export function extractMermaidFromMarkdown(response) {
       fullMatch: match[0],
       startIndex: match.index,
       endIndex: match.index + match[0].length,
-      attributes: attributes
+      attributes: attributes,
+      isInJson: false
     });
   }
 
@@ -868,9 +1019,84 @@ export function extractMermaidFromMarkdown(response) {
   return { diagrams, cleanedResponse: response };
 }
 
+/**
+ * Replace mermaid diagrams in JSON string values with corrected versions
+ * @param {string} originalResponse - Original response with JSON
+ * @param {Array} correctedDiagrams - Array of corrected diagram objects with jsonPath
+ * @returns {string} - Response with corrected diagrams properly escaped in JSON
+ */
+export function replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams) {
+  if (!originalResponse || typeof originalResponse !== 'string') {
+    return originalResponse;
+  }
+
+  if (!correctedDiagrams || correctedDiagrams.length === 0) {
+    return originalResponse;
+  }
+
+  // Extract and parse JSON
+  const jsonResult = extractMermaidFromJson(originalResponse);
+  if (!jsonResult.parsedJson) {
+    return originalResponse;
+  }
+
+  let modifiedJson = jsonResult.parsedJson;
+
+  // Replace diagrams in the JSON object
+  for (const diagram of correctedDiagrams) {
+    if (!diagram.jsonPath || !diagram.isInJson) {
+      continue;
+    }
+
+    // Navigate to the path and replace the content
+    const pathParts = diagram.jsonPath.split('.').filter(p => p);
+    let current = modifiedJson;
+
+    for (let i = 0; i < pathParts.length - 1; i++) {
+      const part = pathParts[i];
+      if (part.startsWith('[') && part.endsWith(']')) {
+        const index = parseInt(part.slice(1, -1), 10);
+        current = current[index];
+      } else {
+        current = current[part];
+      }
+    }
+
+    // Get the last key/index
+    const lastPart = pathParts[pathParts.length - 1];
+    const attributesStr = diagram.attributes ? ` ${diagram.attributes}` : '';
+    const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${diagram.content}\n\`\`\``;
+
+    if (lastPart.startsWith('[') && lastPart.endsWith(']')) {
+      const index = parseInt(lastPart.slice(1, -1), 10);
+      const originalString = current[index];
+      // The fullMatch from extraction has unescaped newlines, so we need to match that
+      current[index] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    } else {
+      const originalString = current[lastPart];
+      // The fullMatch from extraction has unescaped newlines, so we need to match that
+      current[lastPart] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    }
+  }
+
+  // Reconstruct the response with modified JSON
+  const modifiedJsonString = JSON.stringify(modifiedJson, null, 2);
+
+  // Check if original was in a code block
+  const trimmed = originalResponse.trim();
+  if (trimmed.match(/```json\s*\n([\s\S]*?)\n```/)) {
+    return originalResponse.replace(/```json\s*\n([\s\S]*?)\n```/, `\`\`\`json\n${modifiedJsonString}\n\`\`\``);
+  } else if (trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/)) {
+    return originalResponse.replace(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/, `\`\`\`\n${modifiedJsonString}\n\`\`\``);
+  }
+
+  return modifiedJsonString;
+}
+
 /**
  * Replace mermaid diagrams in original markdown with corrected versions
- * @param {string} originalResponse - Original response with markdown
+ * Automatically detects JSON vs markdown format and uses appropriate replacement
+ * @param {string} originalResponse - Original response with markdown or JSON
  * @param {Array} correctedDiagrams - Array of corrected diagram objects
  * @returns {string} - Response with corrected diagrams in original format
  */
@@ -883,22 +1109,28 @@ export function replaceMermaidDiagramsInMarkdown(originalResponse, correctedDiag
     return originalResponse;
   }
 
+  // Check if any diagrams are in JSON format
+  const hasJsonDiagrams = correctedDiagrams.some(d => d.isInJson);
+  if (hasJsonDiagrams) {
+    return replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams);
+  }
+
   let modifiedResponse = originalResponse;
-  
+
   // Sort diagrams by start index in reverse order to preserve indices during replacement
   const sortedDiagrams = [...correctedDiagrams].sort((a, b) => b.startIndex - a.startIndex);
-  
+
   for (const diagram of sortedDiagrams) {
     // Reconstruct the code block with original attributes if they existed
     const attributesStr = diagram.attributes ? ` ${diagram.attributes}` : '';
     const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${diagram.content}\n\`\`\``;
-    
+
     // Replace the original code block
-    modifiedResponse = modifiedResponse.slice(0, diagram.startIndex) + 
-                     newCodeBlock + 
+    modifiedResponse = modifiedResponse.slice(0, diagram.startIndex) +
+                     newCodeBlock +
                      modifiedResponse.slice(diagram.endIndex);
   }
-  
+
   return modifiedResponse;
 }