From d99bd1e0a98f9c81311f02eb9720c2fbf916d591 Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Thu, 13 Feb 2025 18:14:58 -0800
Subject: [PATCH] Add more thorough check for JSON in providers/

---
 providers/service.ts |  9 +++++----
 providers/utils.ts   | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/providers/service.ts b/providers/service.ts
index e876a28ef2..c0a070c455 100644
--- a/providers/service.ts
+++ b/providers/service.ts
@@ -1,6 +1,7 @@
 import { executeTool, getTool } from '@/tools'
 import { getProvider } from './registry'
 import { ProviderRequest, ProviderResponse, TokenInfo } from './types'
+import { extractAndParseJSON } from './utils'
 
 // Helper function to generate provider-specific structured output instructions
 function generateStructuredOutputInstructions(responseFormat: any): string {
@@ -101,8 +102,8 @@ export async function executeProviderRequest(
       // If responseFormat is specified and we have content (not a function call), validate and parse the response
       if (request.responseFormat && content && !provider.hasFunctionCall(currentResponse)) {
         try {
-          // Try to parse the content as JSON
-          const parsedContent = JSON.parse(content)
+          // Extract and parse the JSON content
+          const parsedContent = extractAndParseJSON(content)
 
           // Validate that all required fields are present and have correct types
           const validationErrors = request.responseFormat.fields
@@ -138,8 +139,8 @@ export async function executeProviderRequest(
           // Store the validated JSON response
           content = JSON.stringify(parsedContent)
         } catch (error: any) {
-          console.error('Error parsing structured response:', error)
-          throw new Error(`Failed to parse response as structured output: ${error.message}`)
+          console.error('Raw content:', content)
+          throw new Error(`Failed to parse structured response: ${error.message}`)
         }
       }
 
diff --git a/providers/utils.ts b/providers/utils.ts
index 1bc9d9095f..c646e18def 100644
--- a/providers/utils.ts
+++ b/providers/utils.ts
@@ -38,3 +38,40 @@ export function getProviderFromModel(model: string): ProviderId {
   // Default to deepseek for any other models
   return 'deepseek'
 }
+
+export function extractAndParseJSON(content: string): any {
+  // First clean up the string
+  const trimmed = content.trim()
+
+  // Find the first '{' and last '}'
+  const firstBrace = trimmed.indexOf('{')
+  const lastBrace = trimmed.lastIndexOf('}')
+
+  if (firstBrace === -1 || lastBrace === -1) {
+    throw new Error('No JSON object found in content')
+  }
+
+  // Extract just the JSON part
+  const jsonStr = trimmed.slice(firstBrace, lastBrace + 1)
+
+  try {
+    return JSON.parse(jsonStr)
+  } catch (error) {
+    // If parsing fails, try to clean up common issues
+    const cleaned = jsonStr
+      .replace(/\n/g, ' ') // Remove newlines
+      .replace(/\s+/g, ' ') // Normalize whitespace
+      .replace(/,\s*([}\]])/g, '$1') // Remove trailing commas
+
+    try {
+      return JSON.parse(cleaned)
+    } catch (innerError) {
+      console.error('Original content:', content)
+      console.error('Extracted JSON:', jsonStr)
+      console.error('Cleaned JSON:', cleaned)
+      throw new Error(
+        `Failed to parse JSON after cleanup: ${innerError instanceof Error ? innerError.message : 'Unknown error'}`
+      )
+    }
+  }
+}