Update app.py

app.py

@@ -241,7 +241,7 @@ class ClinicalKPIMonitoring():
 class DiagnosisSupport(ABC):
   """Abstract class for implementing clinical diagnoses."""
   @abstractmethod
-  def diagnose(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
+  def diagnose(self, data: pd.DataFrame, target_col: str, columns: List[str], diagnosis_key : str = "diagnosis" , **kwargs) -> pd.DataFrame:
          pass
 
 class SimpleDiagnosis(DiagnosisSupport):
@@ -249,35 +249,35 @@ class SimpleDiagnosis(DiagnosisSupport):
     def __init__(self):
       self.model : LogisticRegressionTrainer = LogisticRegressionTrainer()
     
-    def diagnose(self, data: pd.DataFrame, target_col: str, columns: List[str], **kwargs) -> pd.DataFrame:
+    def diagnose(self, data: pd.DataFrame, target_col: str, columns: List[str], diagnosis_key : str = "diagnosis", **kwargs) -> pd.DataFrame:
         try:
              result = self.model.invoke(data, target_col=target_col, columns = columns)
              if "accuracy" in result:
-               return pd.DataFrame({"diagnosis": [f"Accuracy {result['accuracy']}"],
+               return pd.DataFrame({diagnosis_key: [f"Accuracy {result['accuracy']}"],
                                    "model": result["model_type"]})
              else:
-                return pd.DataFrame({"diagnosis": [f"Diagnosis failed: {result}"]})
+                return pd.DataFrame({diagnosis_key: [f"Diagnosis failed: {result}"]})
 
         except Exception as e:
-                return pd.DataFrame({"diagnosis":[f"Error during diagnosis {e}"]})
+                return pd.DataFrame({diagnosis_key:[f"Error during diagnosis {e}"]})
 
 
 class TreatmentRecommendation(ABC):
   """Abstract class for treatment recommendations"""
   @abstractmethod
-  def recommend(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
+  def recommend(self, data: pd.DataFrame,  condition_col: str, treatment_col:str, recommendation_key: str = "recommendation", **kwargs) -> pd.DataFrame:
      pass
 
 class BasicTreatmentRecommendation(TreatmentRecommendation):
   """A placeholder class for basic treatment recommendations"""
-  def recommend(self, data: pd.DataFrame, condition_col: str, treatment_col:str, **kwargs) -> pd.DataFrame:
+  def recommend(self, data: pd.DataFrame, condition_col: str, treatment_col:str, recommendation_key: str = "recommendation", **kwargs) -> pd.DataFrame:
        if condition_col not in data.columns or treatment_col not in data.columns:
-          return pd.DataFrame({"recommendation": ["Condition or Treatment columns not found!"]})
+          return pd.DataFrame({recommendation_key: ["Condition or Treatment columns not found!"]})
        treatment = data[data[condition_col] == "High"][treatment_col].to_list()
        if len(treatment)>0:
-         return pd.DataFrame({"recommendation": [f"Treatment recommended for High risk patients: {treatment}"]})
+         return pd.DataFrame({recommendation_key: [f"Treatment recommended for High risk patients: {treatment}"]})
        else:
-           return pd.DataFrame({"recommendation": [f"No treatment recommendation found!"]})
+           return pd.DataFrame({recommendation_key: [f"No treatment recommendation found!"]})
 
 
 class MedicalKnowledgeBase():
@@ -289,12 +289,15 @@ class MedicalKnowledgeBase():
 class SimpleMedicalKnowledge(MedicalKnowledgeBase):
     """Simple Medical Knowledge Class"""
     def search_medical_info(self, query: str) -> str:
+        try:
           if "diabetes treatment" in query.lower():
             return "The recommended treatment for diabetes includes lifestyle changes, medication, and monitoring"
           elif "heart disease risk factors" in query.lower():
                return "Risk factors for heart disease include high blood pressure, high cholesterol, and smoking"
           else:
-            return "No specific information is available"
+              return "No specific information is available"
+        except Exception as e:
+          return f"Medical Knowledge Search Failed {e}"
 
 
 class ForecastingEngine(ABC):
@@ -665,27 +668,51 @@ def main():
                           st.json(result)
             with insights_tab:
                 if selected_data_key:
-                  data = st.session_state.data[selected_data_key]
-                  available_analysis = ["EDA", "temporal", "distribution", "hypothesis", "model"]
-                  selected_analysis = st.multiselect("Select Analysis", available_analysis)
-                  if st.button("Generate Automated Insights"):
-                    with st.spinner("Generating Insights"):
-                        results = st.session_state.automated_insights.generate_insights(data, analysis_names=selected_analysis)
-                        st.json(results)
-                  st.subheader("Diagnosis Support")
-                  target_col = st.selectbox("Select Target Variable for Diagnosis", data.columns.tolist())
-                  num_cols = data.select_dtypes(include=np.number).columns.tolist()
-                  selected_cols_diagnosis = st.multiselect("Select Feature Variables for Diagnosis", num_cols)
-                  if st.button("Generate Diagnosis"):
-                      if target_col and selected_cols_diagnosis:
-                          with st.spinner("Generating Diagnosis"):
-                               result = st.session_state.diagnosis_support.diagnose(data, target_col=target_col, columns=selected_cols_diagnosis)
-                               st.json(result)
-
-                  st.subheader("Treatment Recommendation")
-                  condition_col = st.selectbox("Select Condition Column for Treatment Recommendation", data.columns.tolist())
-                  treatment_col = st.selectbox("Select Treatment Column for Treatment Recommendation", data.columns.tolist())
-                  if st.button("Generate Treatment Recommendation"):
-                      if condition_col and treatment_col:
-                          with st.spinner("Generating Treatment Recommendation"):
-                              result = st.session_state.treatment_recommendation.recommend(data, condition_col = condition_col
+                    data = st.session_state.data[selected_data_key]
+                    available_analysis = ["EDA", "temporal", "distribution", "hypothesis", "model"]
+                    selected_analysis = st.multiselect("Select Analysis", available_analysis)
+                    if st.button("Generate Automated Insights"):
+                        with st.spinner("Generating Insights"):
+                            results = st.session_state.automated_insights.generate_insights(data, analysis_names=selected_analysis)
+                            st.json(results)
+                    st.subheader("Diagnosis Support")
+                    target_col = st.selectbox("Select Target Variable for Diagnosis", data.columns.tolist())
+                    num_cols = data.select_dtypes(include=np.number).columns.tolist()
+                    selected_cols_diagnosis = st.multiselect("Select Feature Variables for Diagnosis", num_cols)
+                    if st.button("Generate Diagnosis"):
+                        if target_col and selected_cols_diagnosis:
+                            with st.spinner("Generating Diagnosis"):
+                                result = st.session_state.diagnosis_support.diagnose(data, target_col=target_col, columns=selected_cols_diagnosis, diagnosis_key="diagnosis_result")
+                                st.json(result)
+
+                    st.subheader("Treatment Recommendation")
+                    condition_col = st.selectbox("Select Condition Column for Treatment Recommendation", data.columns.tolist())
+                    treatment_col = st.selectbox("Select Treatment Column for Treatment Recommendation", data.columns.tolist())
+                    if st.button("Generate Treatment Recommendation"):
+                        if condition_col and treatment_col:
+                            with st.spinner("Generating Treatment Recommendation"):
+                                result = st.session_state.treatment_recommendation.recommend(data, condition_col = condition_col, treatment_col = treatment_col, recommendation_key="treatment_recommendation")
+                                st.json(result)
+
+            with reports_tab:
+                 st.header("Reports")
+                 report_name = st.text_input("Report Name")
+                 report_def = st.text_area("Report definition")
+                 if st.button("Create Report Definition"):
+                   st.session_state.automated_reports.create_report_definition(report_name, report_def)
+                   st.success("Report definition created")
+                 if selected_data_key:
+                   data = st.session_state.data
+                   if st.button("Generate Report"):
+                     with st.spinner("Generating Report..."):
+                       report = st.session_state.automated_reports.generate_report(report_name, data)
+            with knowledge_tab:
+              st.header("Medical Knowledge")
+              query = st.text_input("Enter your medical question here:")
+              if st.button("Search"):
+                  with st.spinner("Searching..."):
+                      result = st.session_state.knowledge_base.search_medical_info(query)
+                      st.write(result)
+
+if __name__ == "__main__":
+    main()