Update hedis_engine.py

hedis_engine.py

@@ -1,4 +1,5 @@
 # hedis_engine.py
+# hedis_engine.py - Updated with Interactive Features
 import os
 import json
 from crewai import Agent, Task, Crew, Process
@@ -7,11 +8,13 @@ from langchain_community.vectorstores import Chroma
 from langchain_openai import ChatOpenAI
 
 
-class InteractiveHedisComplianceEngine:
+class HedisComplianceEngine:
     """
-    Interactive Crew that: (1) states official HEDIS measure criteria and gets user feedback,
+    Crew that: (1) states official HEDIS measure criteria,
     (2) extracts patient evidence from embeddings, and
     (3) produces a markdown analysis report.
+    
+    Now includes interactive feedback capabilities.
     """
     def __init__(self, vectordb: Chroma, measure_code: str, measurement_year: int):
         self.vectordb = vectordb
@@ -144,10 +147,10 @@ CRITICAL FORMATTING REQUIREMENTS:
             # Remove any leading/trailing whitespace and newlines
             result_str = result_str.strip()
             
-            # Remove surrounding quotes if the entire JSON is quoted
+            # Handle quoted JSON strings
             if len(result_str) >= 2 and result_str.startswith('"') and result_str.endswith('"'):
-                # Check if it's actually a quoted JSON string by trying to parse the unquoted version
                 try:
+                    # Try to unquote and parse
                     unquoted = result_str[1:-1]
                     # Handle escaped quotes in the JSON
                     unquoted = unquoted.replace('\\"', '"').replace('\\n', '\n')
@@ -160,6 +163,18 @@ CRITICAL FORMATTING REQUIREMENTS:
             # Try to parse as JSON
             criteria_json = json.loads(result_str)
             return criteria_json
+            
+        except json.JSONDecodeError as e:
+            print(f"Primary method failed, trying simplified approach...")
+            try:
+                return self.get_hedis_criteria_simple(suggestions)
+            except Exception as fallback_error:
+                print(f"Fallback method also failed: {fallback_error}")
+                return {"raw_output": str(result), "parse_error": f"Primary: {str(e)}, Fallback: {str(fallback_error)}"}
+        except Exception as e:
+            print(f"Unexpected error: {e}")
+            return {"raw_output": str(result) if 'result' in locals() else "No result", "parse_error": str(e)}
+
     def get_hedis_criteria_simple(self, suggestions: str = "") -> dict:
         """
         Simplified method to get HEDIS criteria with minimal string formatting.
@@ -214,12 +229,6 @@ CRITICAL FORMATTING REQUIREMENTS:
         except Exception as e:
             print(f"Simple method error: {e}")
             return {"raw_output": str(result) if 'result' in locals() else "No result", "parse_error": str(e)}
-            print(f"JSON parsing error: {e}")
-            print(f"Raw result type: {type(result)}")
-            print(f"Raw result: {repr(str(result))}")
-            print(f"Cleaned result: {repr(result_str if 'result_str' in locals() else 'N/A')}")
-            # If not valid JSON, return as string
-            return {"raw_output": str(result), "parse_error": f"JSON parsing failed: {str(e)}"}
 
     def display_criteria(self, criteria: dict) -> None:
         """Display the HEDIS criteria in a user-friendly format."""
@@ -396,19 +405,21 @@ CRITICAL FORMATTING REQUIREMENTS:
         print(f"\n✅ Analysis complete!")
         return str(result)
 
-    def run_non_interactive(self) -> str:
+    def run(self) -> str:
         """
-        Run the original non-interactive analysis (for backwards compatibility).
+        Entry point - runs original non-interactive analysis for app compatibility.
+        Use run_interactive() for interactive mode.
         """
+        # Original tasks
         hedis_task = Task(
             description=(
-                f"Provide the official criteria for HEDIS measure {self.measure_code}, including:\n"
+                "Provide the official criteria for HEDIS measure {measure_code}, including:\n"
                 "- Required tests or procedures\n"
                 "- Diagnosis/procedure codes\n"
                 "- Required medications\n"
                 "- Inclusion and exclusion criteria\n"
-                f"- Timeframes for each requirement (e.g., within X years of {self.measurement_year})\n\n"
-                f"You MUST compute and include absolute date ranges using {self.measurement_year}.\n"
+                "- Timeframes for each requirement (e.g., within 10 years of {measurement_year})\n\n"
+                "You MUST compute and include absolute date ranges using {measurement_year}.\n"
                 "Return JSON with fields: 'measure','required_tests','required_medications',"
                 "'codes','timeframes','inclusions','exclusions'."
             ),
@@ -419,8 +430,8 @@ CRITICAL FORMATTING REQUIREMENTS:
 
         patient_task = Task(
             description=(
-                f"Using the patient_chart_search tool, extract only the tests, diagnoses, procedures, "
-                f"medications, and dates relevant to {self.measure_code}. Only include measure related info from the patient charts"
+                "Using the patient_chart_search tool, extract only the tests, diagnoses, procedures, "
+                "medications, and dates relevant to {measure_code}. Only include measure related info from the patient charts"
             ),
             expected_output="Structured list of patient evidence with dates.",
             agent=self.patient_analyzer,
@@ -429,8 +440,8 @@ CRITICAL FORMATTING REQUIREMENTS:
 
         chart_analysis_task = Task(
             description=(
-                f"Create a comprehensive markdown report analyzing the patient chart for HEDIS {self.measure_code} "
-                f"for measurement year {self.measurement_year}.\n\n"
+                "Create a comprehensive markdown report analyzing the patient chart for HEDIS {measure_code} "
+                "for measurement year {measurement_year}.\n\n"
                 "Sections:\n"
                 "### Diagnoses Found\n"
                 "### Procedures Found\n"
@@ -459,27 +470,4 @@ CRITICAL FORMATTING REQUIREMENTS:
             "measurement_year": self.measurement_year,
         }
         result = crew.kickoff(inputs=inputs)
-        return str(result)
-
-    def run(self) -> str:
-        """
-        Entry point - runs non-interactive analysis by default for app compatibility.
-        Use run_interactive() for interactive mode.
-        """
-        return self.run_non_interactive()
-
-
-# Usage example:
-"""
-# Initialize the engine
-vectordb = your_chroma_db  # Your Chroma vector database
-engine = InteractiveHedisComplianceEngine(
-    vectordb=vectordb,
-    measure_code="CDC-A1C",
-    measurement_year=2024
-)
-
-# Run interactive analysis
-result = engine.run()
-print(result)
-"""
+        return str(result)