Update TestFindingAgent.py

TestFindingAgent.py

@@ -3,14 +3,13 @@ import json
 import pandas as pd
 from PyPDF2 import PdfReader
 from json_repair import repair_json
-from typing import List, Dict, Any
 from typing import List, Dict, Any, Optional
 from crewai import Agent, Task, Crew, Process
 from crewai_tools import SerperDevTool
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_openai import ChatOpenAI
 from langchain_community.vectorstores import Chroma
 
+
 SEED_SOURCES = [
     "https://www.cms.gov/medicare/payment/medicare-advantage-rates-statistics/risk-adjustment",
     "https://www.cms.gov/data-research/monitoring-programs/medicare-risk-adjustment-data-validation-program",
@@ -45,6 +44,30 @@ class TestFindingAgent:
             llm=self.llm,
         )
 
+    def _extract_json_from_llm(self, raw_response: str) -> Dict[str, Any]:
+        """Extracts and repairs JSON from an LLM response safely."""
+        import re
+        match = re.search(r"\{.*\}", raw_response, re.DOTALL)
+        if not match:
+            print("[ERROR] No JSON object found in LLM response")
+            return {}
+
+        clean_json_str = match.group(0)
+
+        # Step 1: Try direct JSON parse
+        try:
+            return json.loads(clean_json_str)
+        except json.JSONDecodeError as e:
+            print(f"[WARN] Direct JSON parsing failed: {e}")
+
+        # Step 2: Try repairing JSON
+        try:
+            repaired = repair_json(clean_json_str)
+            return json.loads(repaired)
+        except Exception as e:
+            print(f"[ERROR] Failed to repair and parse JSON: {e}")
+            return {}
+
     def run(self, input_diagnoses: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         updated_list = []
 
@@ -77,13 +100,11 @@ class TestFindingAgent:
                 verbose=True
             )
             result = crew.kickoff()
-            result = json.loads(repair_json(result))
-            diag["tests"] = result
-            updated_list.append(diag)
 
-        # Save results to file
-        # with open(self.output_file, "w", encoding="utf-8") as f:
-        #     json.dump(updated_list, f, indent=2, ensure_ascii=False)
+            # Use safe extractor
+            result_dict = self._extract_json_from_llm(result)
+
+            diag["tests"] = result_dict
+            updated_list.append(diag)
 
-        # print(f"[OUTPUT] Saved {len(updated_list)} diagnoses with appended tests to {self.output_file}")
-        return updated_list
+        return updated_list