Update app.py

app.py

@@ -10,6 +10,7 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
 import streamlit as st
+import spacy
 
 from scipy.stats import ttest_ind, f_oneway
 from sklearn.model_selection import train_test_split
@@ -22,7 +23,7 @@ from statsmodels.tsa.seasonal import seasonal_decompose
 from statsmodels.tsa.stattools import adfuller
 
 from pydantic import BaseModel, Field
-from Bio import Entrez
+from Bio import Entrez  # Ensure BioPython is installed
 
 from langchain.prompts import PromptTemplate
 from groq import Groq
@@ -31,6 +32,14 @@ from groq import Groq
 # Initialize Groq Client with API Key from environment variables
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
+# Load spaCy model with error handling
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    st.write("Downloading en_core_web_sm spaCy model...")
+    spacy.cli.download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
+
 # ---------------------- Base Classes and Schemas ---------------------------
 
 class ResearchInput(BaseModel):
@@ -68,7 +77,7 @@ class AdvancedEDA(DataAnalyzer):
                 "dimensionality": {
                     "rows": len(data),
                     "columns": list(data.columns),
-                    "memory_usage": f"{data.memory_usage().sum() / 1e6:.2f} MB"
+                    "memory_usage_MB": f"{data.memory_usage().sum() / 1e6:.2f} MB"
                 },
                 "statistical_profile": data.describe(percentiles=[.25, .5, .75]).to_dict(),
                 "temporal_analysis": {
@@ -127,10 +136,12 @@ class TemporalAnalyzer(DataAnalyzer):
             plt.close()
             plot_data = base64.b64encode(buf.getvalue()).decode()
 
+            stationarity_p_value = adfuller(ts_data)[1]
+
             return {
                 "trend_statistics": {
-                    "stationarity_p_value": adfuller(ts_data)[1],
-                    "seasonality_strength": max(decomposition.seasonal)
+                    "stationarity_p_value": stationarity_p_value,
+                    "seasonality_strength": float(max(decomposition.seasonal))
                 },
                 "visualization": plot_data
             }
@@ -197,7 +208,9 @@ class LogisticRegressionTrainer(DataAnalyzer):
         try:
             X = data[columns]
             y = data[target_col]
-            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+            X_train, X_test, y_train, y_test = train_test_split(
+                X, y, test_size=0.2, random_state=42
+            )
             model = LogisticRegression(max_iter=1000)
             model.fit(X_train, y_train)
             y_pred = model.predict(X_test)
@@ -232,7 +245,8 @@ class ClinicalRulesEngine:
         for rule_name, rule in self.rules.items():
             try:
                 # Evaluate the condition using the dataframe 'df'
-                rule_matched = eval(rule.condition, {}, {"df": data})
+                # **Warning**: Using eval can be dangerous. Ensure that user inputs are sanitized.
+                rule_matched = eval(rule.condition, {"__builtins__": None}, {"df": data})
                 results[rule_name] = {
                     "rule_matched": rule_matched,
                     "action": rule.action if rule_matched else None,
@@ -264,7 +278,8 @@ class ClinicalKPIMonitoring:
         results = {}
         for kpi_name, kpi in self.kpis.items():
             try:
-                kpi_value = eval(kpi.calculation, {}, {"df": data})
+                # **Warning**: Using eval can be dangerous. Ensure that user inputs are sanitized.
+                kpi_value = eval(kpi.calculation, {"__builtins__": None}, {"df": data})
                 results[kpi_name] = {
                     "value": kpi_value,
                     "threshold": kpi.threshold,
@@ -363,6 +378,8 @@ class BasicTreatmentRecommendation(TreatmentRecommendation):
                 recommendation_key: ["No treatment recommendation found!"]
             })
 
+# ---------------------- Medical Knowledge Base ---------------------------
+
 class MedicalKnowledgeBase(ABC):
     """Abstract class for Medical Knowledge."""
     @abstractmethod
@@ -370,14 +387,15 @@ class MedicalKnowledgeBase(ABC):
         pass
 
 class SimpleMedicalKnowledge(MedicalKnowledgeBase):
-    """Simple Medical Knowledge Class with TF-IDF and PubMed."""
+    """Simple Medical Knowledge Class with TF-IDF, NER, and PubMed."""
     def __init__(self):
         self.knowledge_base = {
             "diabetes": "The recommended treatment for diabetes includes lifestyle changes, medication, and monitoring.",
             "heart disease": "Risk factors for heart disease include high blood pressure, high cholesterol, and smoking.",
             "fever": "For a fever, you can consider over-the-counter medications like acetaminophen or ibuprofen. Rest and hydration are also important.",
             "headache": "For a headache, try rest, hydration, and over-the-counter pain relievers. Consult a doctor if it is severe or persistent.",
-            "cold": "For a cold, get rest, drink plenty of fluids, and use over-the-counter remedies like decongestants."
+            "cold": "For a cold, get rest, drink plenty of fluids, and use over-the-counter remedies like decongestants.",
+            "cancer drugs": "Please consult with your doctor to get personalized treatment and the latest drug information for your type of cancer."
         }
         self.vectorizer = TfidfVectorizer()
         self.tfidf_matrix = self.vectorizer.fit_transform(self.knowledge_base.values())
@@ -402,26 +420,61 @@ class SimpleMedicalKnowledge(MedicalKnowledgeBase):
     def search_medical_info(self, query: str, pub_email: str = "") -> str:
         """Search the medical knowledge base and PubMed for relevant information."""
         try:
-            query_vector = self.vectorizer.transform([query])
-            similarities = cosine_similarity(query_vector, self.tfidf_matrix)
-            best_match_index = np.argmax(similarities)
-            best_match_keyword = list(self.knowledge_base.keys())[best_match_index]
-            best_match_info = list(self.knowledge_base.values())[best_match_index]
-
-            pubmed_result = self.search_pubmed(query, pub_email)
-            if "No abstracts found" not in pubmed_result:
-                return (
-                    f"**Based on your query:** {best_match_info}\n\n"
-                    f"**PubMed Abstract:**\n{pubmed_result}"
-                )
+            query_lower = query.lower()
+            doc = nlp(query_lower)
+            entities = [ent.text for ent in doc.ents]
+
+            if entities:
+                best_match_keyword = ""
+                best_match_score = -1
+                for entity in entities:
+                    query_vector = self.vectorizer.transform([entity])
+                    similarities = cosine_similarity(query_vector, self.tfidf_matrix)
+                    current_best_match_index = np.argmax(similarities)
+                    current_best_score = np.max(similarities)
+                    if current_best_score > best_match_score:
+                        best_match_keyword = list(self.knowledge_base.keys())[current_best_match_index]
+                        best_match_score = current_best_score
             else:
-                return (
-                    f"**Based on your query:** {best_match_info}\n\n"
-                    f"{pubmed_result}"
-                )
+                query_vector = self.vectorizer.transform([query_lower])
+                similarities = cosine_similarity(query_vector, self.tfidf_matrix)
+                best_match_index = np.argmax(similarities)
+                best_match_keyword = list(self.knowledge_base.keys())[best_match_index]
+
+            best_match_info = self.knowledge_base.get(best_match_keyword, "No specific information is available based on the query provided.")
+
+            pubmed_result = self.search_pubmed(best_match_keyword, pub_email)
+            feedback_key = f"feedback_{query_lower}"  # Unique key for feedback
+
+            response = f"**Based on your query:** {best_match_info}\n\n"
+
+            if "No abstracts found for this query on PubMed" not in pubmed_result:
+                response += f"**PubMed Abstract:**\n{pubmed_result}"
+            else:
+                response += f"{pubmed_result}"
+
+            # Initialize feedback in session state
+            if feedback_key not in st.session_state:
+                st.session_state[feedback_key] = {"feedback": None}
+
+            # Display feedback buttons only if a valid response is generated
+            if "error" not in pubmed_result:
+                col1, col2 = st.columns([1, 1])
+                with col1:
+                    if st.button("Good Result", key=f"good_{feedback_key}"):
+                        st.session_state[feedback_key]["feedback"] = "positive"
+                        st.success("Thank you for the feedback!")
+                with col2:
+                    if st.button("Bad Result", key=f"bad_{feedback_key}"):
+                        st.session_state[feedback_key]["feedback"] = "negative"
+                        st.error("Thank you for the feedback!")
+            
+            return response
         except Exception as e:
             return f"Medical Knowledge Search Failed: {e}"
 
+# ---------------------- Forecasting Engine ---------------------------
+
 class ForecastingEngine(ABC):
     """Abstract class for forecasting."""
     @abstractmethod
@@ -472,23 +525,19 @@ class Dashboard:
         st.header("Dashboard")
         for vis_name, vis_type in self.layout.items():
             st.subheader(vis_name)
-            if vis_type == "table":
-                df = data_dict.get(vis_name)
-                if df is not None:
+            df = data_dict.get(vis_name)
+            if df is not None:
+                if vis_type == "table":
                     st.table(df)
-                else:
-                    st.write("Data Not Found")
-            elif vis_type == "plot":
-                df = data_dict.get(vis_name)
-                if df is not None:
+                elif vis_type == "plot":
                     if len(df.columns) > 1:
                         fig = plt.figure()
                         sns.lineplot(data=df)
                         st.pyplot(fig)
                     else:
                         st.write("Please select a DataFrame with more than 1 column for plotting.")
-                else:
-                    st.write("Data not found")
+            else:
+                st.write("Data Not Found")
 
 class AutomatedReports:
     """Manages automated report definitions and generation."""
@@ -754,9 +803,9 @@ def main_tabs_section():
 
 def data_analysis_section():
     """Handles the Data Analysis tab."""
-    selected_data_key = st.sidebar.selectbox("Select Dataset for Analysis", list(st.session_state.data.keys()))
+    selected_data_key = st.session_state.get('selected_data_key', None)
     if not selected_data_key:
-        st.warning("Please select a dataset to perform analysis.")
+        st.warning("Please select a dataset from the metadata section.")
         return
 
     data = st.session_state.data[selected_data_key]
@@ -791,21 +840,33 @@ def perform_temporal_analysis(data: pd.DataFrame):
     time_cols = data.select_dtypes(include='datetime').columns
     num_cols = data.select_dtypes(include=np.number).columns
 
+    if len(time_cols) == 0:
+        st.warning("No datetime columns available for temporal analysis.")
+        return
+
     time_col = st.selectbox("Select Temporal Variable", time_cols)
     value_col = st.selectbox("Select Analysis Variable", num_cols)
 
     if time_col and value_col:
         analyzer = TemporalAnalyzer()
         result = analyzer.invoke(data=data, time_col=time_col, value_col=value_col)
-        if "visualization" in result:
-            st.image(f"data:image/png;base64,{result['visualization']}")
+        if "visualization" in result and result["visualization"]:
+            st.image(f"data:image/png;base64,{result['visualization']}", use_column_width=True)
         st.json(result)
 
 def perform_comparative_statistics(data: pd.DataFrame):
     """Performs Comparative Statistics."""
-    categorical_cols = data.select_dtypes(include='category').columns
+    categorical_cols = data.select_dtypes(include=['category', 'object']).columns
     numeric_cols = data.select_dtypes(include=np.number).columns
 
+    if len(categorical_cols) == 0:
+        st.warning("No categorical columns available for hypothesis testing.")
+        return
+
+    if len(numeric_cols) == 0:
+        st.warning("No numerical columns available for hypothesis testing.")
+        return
+
     group_col = st.selectbox("Select Grouping Variable", categorical_cols)
     value_col = st.selectbox("Select Metric Variable", numeric_cols)
 
@@ -823,10 +884,12 @@ def perform_distribution_analysis(data: pd.DataFrame):
     if selected_cols:
         analyzer = DistributionVisualizer()
         img_data = analyzer.invoke(data=data, columns=selected_cols)
-        if "Visualization Error" not in img_data:
-            st.image(f"data:image/png;base64,{img_data}")
+        if not img_data.startswith("Visualization Error"):
+            st.image(f"data:image/png;base64,{img_data}", use_column_width=True)
         else:
             st.error(img_data)
+    else:
+        st.info("Please select at least one numerical column to visualize.")
 
 def perform_logistic_regression_training(data: pd.DataFrame):
     """Trains a Logistic Regression model."""
@@ -839,6 +902,8 @@ def perform_logistic_regression_training(data: pd.DataFrame):
         result = analyzer.invoke(data=data, target_col=target_col, columns=selected_cols)
         st.subheader("Logistic Regression Model Results")
         st.json(result)
+    else:
+        st.warning("Please select both target and feature variables for model training.")
 
 def clinical_logic_section():
     """Handles the Clinical Logic tab."""
@@ -853,39 +918,45 @@ def clinical_logic_section():
     severity = st.selectbox("Enter Severity for the Rule", ["low", "medium", "high"])
 
     if st.button("Add Clinical Rule"):
-        try:
-            rule = ClinicalRule(
-                name=rule_name, 
-                condition=condition, 
-                action=action, 
-                severity=severity
-            )
-            st.session_state.clinical_rules.add_rule(rule)
-            st.success("Added Clinical Rule successfully.")
-        except Exception as e:
-            st.error(f"Error in rule definition: {e}")
+        if rule_name and condition and action and severity:
+            try:
+                rule = ClinicalRule(
+                    name=rule_name, 
+                    condition=condition, 
+                    action=action, 
+                    severity=severity
+                )
+                st.session_state.clinical_rules.add_rule(rule)
+                st.success("Added Clinical Rule successfully.")
+            except Exception as e:
+                st.error(f"Error in rule definition: {e}")
+        else:
+            st.error("Please fill in all fields to add a clinical rule.")
 
     # Clinical KPI Management
     st.subheader("Clinical KPI Definition")
     kpi_name = st.text_input("Enter KPI Name")
     kpi_calculation = st.text_area("Enter KPI Calculation (use 'df' for DataFrame)", 
                                     help="Example: df['patient_count'].sum()")
-    threshold = st.text_input("Enter Threshold for KPI", help="Optional")
+    threshold = st.text_input("Enter Threshold for KPI (Optional)", help="Leave blank if not applicable")
 
     if st.button("Add Clinical KPI"):
-        try:
-            threshold_value = float(threshold) if threshold else None
-            kpi = ClinicalKPI(
-                name=kpi_name, 
-                calculation=kpi_calculation, 
-                threshold=threshold_value
-            )
-            st.session_state.kpi_monitoring.add_kpi(kpi)
-            st.success(f"Added KPI '{kpi_name}' successfully.")
-        except ValueError:
-            st.error("Threshold must be a numeric value.")
-        except Exception as e:
-            st.error(f"Error creating KPI: {e}")
+        if kpi_name and kpi_calculation:
+            try:
+                threshold_value = float(threshold) if threshold else None
+                kpi = ClinicalKPI(
+                    name=kpi_name, 
+                    calculation=kpi_calculation, 
+                    threshold=threshold_value
+                )
+                st.session_state.kpi_monitoring.add_kpi(kpi)
+                st.success(f"Added KPI '{kpi_name}' successfully.")
+            except ValueError:
+                st.error("Threshold must be a numeric value.")
+            except Exception as e:
+                st.error(f"Error creating KPI: {e}")
+        else:
+            st.error("Please provide both KPI name and calculation.")
 
     # Execute Clinical Rules and Calculate KPIs
     selected_data_key = st.selectbox("Select Dataset for Clinical Logic", list(st.session_state.data.keys()))
@@ -899,6 +970,8 @@ def clinical_logic_section():
             with st.spinner("Calculating Clinical KPIs..."):
                 result = st.session_state.kpi_monitoring.calculate_kpis(data)
                 st.json(result)
+    else:
+        st.warning("Please ingest data to execute clinical rules and calculate KPIs.")
 
 def insights_section():
     """Handles the Insights tab."""
@@ -914,11 +987,14 @@ def insights_section():
     selected_analyses = st.multiselect("Select Analyses for Insights", available_analyses)
 
     if st.button("Generate Automated Insights"):
-        with st.spinner("Generating Insights..."):
-            results = st.session_state.automated_insights.generate_insights(
-                data, analysis_names=selected_analyses
-            )
-            st.json(results)
+        if selected_analyses:
+            with st.spinner("Generating Insights..."):
+                results = st.session_state.automated_insights.generate_insights(
+                    data, analysis_names=selected_analyses
+                )
+                st.json(results)
+        else:
+            st.warning("Please select at least one analysis to generate insights.")
 
     # Diagnosis Support
     st.subheader("Diagnosis Support")
@@ -958,7 +1034,7 @@ def reports_section():
     # Create Report Definition
     st.subheader("Create Report Definition")
     report_name = st.text_input("Report Name")
-    report_def = st.text_area("Report Definition")
+    report_def = st.text_area("Report Definition", help="Describe the structure and content of the report.")
 
     if st.button("Create Report Definition"):
         if report_name and report_def:
@@ -969,22 +1045,22 @@ def reports_section():
 
     # Generate Report
     st.subheader("Generate Report")
-    report_name_to_generate = st.selectbox("Select Report to Generate", list(st.session_state.automated_reports.report_definitions.keys()))
-    
-    if st.button("Generate Report"):
-        if report_name_to_generate:
+    report_names = list(st.session_state.automated_reports.report_definitions.keys())
+    if report_names:
+        report_name_to_generate = st.selectbox("Select Report to Generate", report_names)
+        if st.button("Generate Report"):
             with st.spinner("Generating Report..."):
                 report = st.session_state.automated_reports.generate_report(report_name_to_generate, st.session_state.data)
                 if "error" not in report:
-                    st.header(f"Report: {report_name_to_generate}")
-                    st.write(f"**Definition:** {report['Report Definition']}")
+                    st.header(f"Report: {report['Report Name']}")
+                    st.markdown(f"**Definition:** {report['Report Definition']}")
                     for df_name, df_content in report["Data"].items():
                         st.subheader(f"Data: {df_name}")
-                        st.write(pd.DataFrame(df_content))
+                        st.dataframe(pd.DataFrame(df_content))
                 else:
                     st.error(report["error"])
-        else:
-            st.error("Please select a report to generate.")
+    else:
+        st.info("No report definitions found. Please create a report definition first.")
 
 def medical_knowledge_section():
     """Handles the Medical Knowledge tab."""
@@ -992,9 +1068,11 @@ def medical_knowledge_section():
     query = st.text_input("Enter your medical question here:")
 
     if st.button("Search"):
-        if query:
+        if query.strip():
             with st.spinner("Searching..."):
-                result = st.session_state.knowledge_base.search_medical_info(query, pub_email=st.session_state.pub_email)
+                result = st.session_state.knowledge_base.search_medical_info(
+                    query, pub_email=st.session_state.pub_email
+                )
                 st.markdown(result)
         else:
             st.error("Please enter a medical question to search.")