From 0ded17d7358ef027900e735adae21604dcb7c08b Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 5 Feb 2025 12:41:20 +0000
Subject: [PATCH] fix: improve JSON parsing and increase test timeouts

Co-Authored-By: Han Xiao <han.xiao@jina.ai>
---
 src/tools/__tests__/dedup.test.ts          | 2 ++
 src/tools/__tests__/error-analyzer.test.ts | 2 ++
 src/tools/__tests__/evaluator.test.ts      | 2 ++
 src/tools/dedup.ts                         | 4 +++-
 src/tools/error-analyzer.ts                | 4 +++-
 src/tools/evaluator.ts                     | 4 +++-
 src/tools/query-rewriter.ts                | 4 +++-
 7 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/tools/__tests__/dedup.test.ts b/src/tools/__tests__/dedup.test.ts
index ed615a7..cd803f9 100644
--- a/src/tools/__tests__/dedup.test.ts
+++ b/src/tools/__tests__/dedup.test.ts
@@ -1,5 +1,7 @@
 import { dedupQueries } from '../dedup';
 
+jest.setTimeout(10000);
+
 describe('dedupQueries', () => {
   it('should remove duplicate queries', async () => {
     const queries = ['typescript tutorial', 'typescript tutorial', 'javascript basics'];
diff --git a/src/tools/__tests__/error-analyzer.test.ts b/src/tools/__tests__/error-analyzer.test.ts
index 7729a0d..481a7b3 100644
--- a/src/tools/__tests__/error-analyzer.test.ts
+++ b/src/tools/__tests__/error-analyzer.test.ts
@@ -1,5 +1,7 @@
 import { analyzeSteps } from '../error-analyzer';
 
+jest.setTimeout(10000);
+
 describe('analyzeSteps', () => {
   it('should analyze error steps', async () => {
     const { response } = await analyzeSteps(['Step 1: Search failed', 'Step 2: Invalid query']);
diff --git a/src/tools/__tests__/evaluator.test.ts b/src/tools/__tests__/evaluator.test.ts
index 391b4bb..9c156ef 100644
--- a/src/tools/__tests__/evaluator.test.ts
+++ b/src/tools/__tests__/evaluator.test.ts
@@ -1,6 +1,8 @@
 import { evaluateAnswer } from '../evaluator';
 import { TokenTracker } from '../../utils/token-tracker';
 
+jest.setTimeout(10000);
+
 describe('evaluateAnswer', () => {
   it('should evaluate answer definitiveness', async () => {
     const tokenTracker = new TokenTracker();
diff --git a/src/tools/dedup.ts b/src/tools/dedup.ts
index 9e5d94a..69ebb0e 100644
--- a/src/tools/dedup.ts
+++ b/src/tools/dedup.ts
@@ -27,8 +27,10 @@ async function generateResponse(provider: AIProvider, prompt: string, providerTy
         }
       });
       const response = await result.response;
+      const text = response.text();
+      const jsonMatch = text.match(/```json\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
       return {
-        text: response.text(),
+        text: jsonMatch ? jsonMatch[1].trim() : text,
         tokens: response.usageMetadata?.totalTokenCount || 0
       };
     }
diff --git a/src/tools/error-analyzer.ts b/src/tools/error-analyzer.ts
index 4be21ef..75ffb89 100644
--- a/src/tools/error-analyzer.ts
+++ b/src/tools/error-analyzer.ts
@@ -114,8 +114,10 @@ async function generateResponse(provider: AIProvider, prompt: string, providerTy
         }
       });
       const response = await result.response;
+      const text = response.text();
+      const jsonMatch = text.match(/```json\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
       return {
-        text: response.text(),
+        text: jsonMatch ? jsonMatch[1].trim() : text,
         tokens: response.usageMetadata?.totalTokenCount || 0
       };
     }
diff --git a/src/tools/evaluator.ts b/src/tools/evaluator.ts
index bf370ea..2fa756d 100644
--- a/src/tools/evaluator.ts
+++ b/src/tools/evaluator.ts
@@ -25,8 +25,10 @@ async function generateResponse(provider: AIProvider, prompt: string, providerTy
         }
       });
       const response = await result.response;
+      const text = response.text();
+      const jsonMatch = text.match(/```json\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
       return {
-        text: response.text(),
+        text: jsonMatch ? jsonMatch[1].trim() : text,
         tokens: response.usageMetadata?.totalTokenCount || 0
       };
     }
diff --git a/src/tools/query-rewriter.ts b/src/tools/query-rewriter.ts
index 42fb338..d85c95a 100644
--- a/src/tools/query-rewriter.ts
+++ b/src/tools/query-rewriter.ts
@@ -105,8 +105,10 @@ async function generateResponse(provider: AIProvider, prompt: string, providerTy
         }
       });
       const response = await result.response;
+      const text = response.text();
+      const jsonMatch = text.match(/```json\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
       return {
-        text: response.text(),
+        text: jsonMatch ? jsonMatch[1].trim() : text,
         tokens: response.usageMetadata?.totalTokenCount || 0
       };
     }