Update app.py

app.py

@@ -60,6 +60,9 @@ def _get_vectordb_from_bytes(pdf_bytes: bytes):
 
 def clean_json_response(response_text: str) -> dict:
     """Clean and parse JSON response from agents."""
+    if not response_text:
+        return {"raw_output": "Empty response", "parse_error": "No content received"}
+        
     try:
         # Clean the response string
         cleaned = str(response_text).strip()
@@ -76,16 +79,26 @@ def clean_json_response(response_text: str) -> dict:
         # Remove any leading/trailing whitespace and newlines
         cleaned = cleaned.strip()
         
-        # Remove any leading/trailing quotes if the entire thing is quoted
-        if cleaned.startswith('"') and cleaned.endswith('"'):
-            cleaned = cleaned[1:-1]
+        # Handle quoted JSON strings
+        if len(cleaned) >= 2 and cleaned.startswith('"') and cleaned.endswith('"'):
+            try:
+                # Try to unquote and parse
+                unquoted = cleaned[1:-1]
+                # Handle common escape sequences
+                unquoted = unquoted.replace('\\"', '"').replace('\\n', '\n').replace('\\t', '\t')
+                test_parse = json.loads(unquoted)
+                cleaned = unquoted
+            except:
+                # If unquoting fails, keep original
+                pass
         
         # Try to parse as JSON
         return json.loads(cleaned)
+        
     except json.JSONDecodeError as e:
         print(f"JSON parsing error: {e}")
-        print(f"Cleaned response: {repr(cleaned)}")
-        return {"raw_output": str(response_text), "parse_error": str(e)}
+        print(f"Original response: {repr(response_text)}")
+        return {"raw_output": str(response_text), "parse_error": f"JSON Parse Error: {str(e)}"}
     """Format criteria dictionary into markdown for display."""
     if "raw_output" in criteria:
         return f"```\n{criteria['raw_output']}\n```"
@@ -202,17 +215,32 @@ def start_hedis_analysis(pdf_file, measure_name, measurement_year):
         # Create engine and get initial criteria
         _CURRENT_ENGINE = InteractiveHedisComplianceEngine(vectordb, measure_name, int(measurement_year))
         
-        # Get initial criteria with better error handling
-        raw_criteria = _CURRENT_ENGINE.get_hedis_criteria_with_feedback()
-        
-        # Clean the criteria if needed
-        if isinstance(raw_criteria, dict) and "raw_output" in raw_criteria:
-            # Try to clean and re-parse if it's raw output
-            _CURRENT_CRITERIA = clean_json_response(raw_criteria["raw_output"])
-        else:
-            _CURRENT_CRITERIA = raw_criteria
-        
-        _ANALYSIS_STATE = "criteria_review"
+        try:
+            # Get initial criteria with better error handling
+            raw_criteria = _CURRENT_ENGINE.get_hedis_criteria_with_feedback()
+            
+            # Handle different response types
+            if isinstance(raw_criteria, dict):
+                if "raw_output" in raw_criteria:
+                    # Try to clean and re-parse if it's raw output
+                    _CURRENT_CRITERIA = clean_json_response(raw_criteria["raw_output"])
+                else:
+                    _CURRENT_CRITERIA = raw_criteria
+            else:
+                # If it's not a dict, treat as raw output
+                _CURRENT_CRITERIA = clean_json_response(str(raw_criteria))
+            
+            _ANALYSIS_STATE = "criteria_review"
+            
+        except Exception as e:
+            print(f"Error getting criteria: {e}")
+            return (
+                f"❌ Error getting HEDIS criteria: {str(e)}\n\nPlease try again or check the measure code.",
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False)
+            )
         
         criteria_display = format_criteria_for_display(_CURRENT_CRITERIA)
         
@@ -292,17 +320,32 @@ def modify_criteria(suggestions):
                 gr.update(visible=True)
             )
         
-        # Get modified criteria with better error handling
-        raw_criteria = _CURRENT_ENGINE.get_hedis_criteria_with_feedback(suggestions.strip())
-        
-        # Clean the criteria if needed
-        if isinstance(raw_criteria, dict) and "raw_output" in raw_criteria:
-            # Try to clean and re-parse if it's raw output
-            _CURRENT_CRITERIA = clean_json_response(raw_criteria["raw_output"])
-        else:
-            _CURRENT_CRITERIA = raw_criteria
-        
-        criteria_display = format_criteria_for_display(_CURRENT_CRITERIA)
+        try:
+            # Get modified criteria with better error handling
+            raw_criteria = _CURRENT_ENGINE.get_hedis_criteria_with_feedback(suggestions.strip())
+            
+            # Handle different response types
+            if isinstance(raw_criteria, dict):
+                if "raw_output" in raw_criteria:
+                    # Try to clean and re-parse if it's raw output
+                    _CURRENT_CRITERIA = clean_json_response(raw_criteria["raw_output"])
+                else:
+                    _CURRENT_CRITERIA = raw_criteria
+            else:
+                # If it's not a dict, treat as raw output
+                _CURRENT_CRITERIA = clean_json_response(str(raw_criteria))
+            
+            criteria_display = format_criteria_for_display(_CURRENT_CRITERIA)
+            
+        except Exception as e:
+            print(f"Error modifying criteria: {e}")
+            return (
+                f"❌ Error modifying criteria: {str(e)}\n\nPlease try again with different suggestions.",
+                gr.update(visible=True, value=""),
+                gr.update(visible=True),
+                gr.update(visible=True),
+                gr.update(visible=True)
+            )
         
         return (
             f"## 🔄 Modified Criteria\n\n{criteria_display}",