Update app.py

app.py

@@ -25,8 +25,12 @@ from statsmodels.tsa.stattools import adfuller
 from pydantic import BaseModel, Field
 from Bio import Entrez  # Ensure BioPython is installed
 
-from langchain.prompts import PromptTemplate
-from groq import Groq
+from dotenv import load_dotenv
+import requests
+import ast
+
+# ---------------------- Load Environment Variables ---------------------------
+load_dotenv()
 
 # ---------------------- Streamlit Page Configuration ---------------------------
 # This must be the first Streamlit command in the script
@@ -34,7 +38,18 @@ st.set_page_config(page_title="AI Clinical Intelligence Hub", layout="wide")
 
 # ---------------------- Initialize External Clients ---------------------------
 # Initialize Groq Client with API Key from environment variables
-client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+GROQ_API_ENDPOINT = os.getenv("GROQ_API_ENDPOINT")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+PUB_EMAIL = os.getenv("PUB_EMAIL", "")
+
+if not GROQ_API_ENDPOINT or not GROQ_API_KEY:
+    st.error("Groq API endpoint and key must be set as environment variables.")
+    st.stop()
+
+headers = {
+    "Authorization": f"Bearer {GROQ_API_KEY}",
+    "Content-Type": "application/json"
+}
 
 # Load spaCy model with error handling
 try:
@@ -250,9 +265,8 @@ class ClinicalRulesEngine:
         results = {}
         for rule_name, rule in self.rules.items():
             try:
-                # **Warning**: Using eval can be dangerous. Ensure that user inputs are sanitized.
-                # Here, we're restricting the environment by removing built-ins.
-                rule_matched = eval(rule.condition, {"__builtins__": None}, {"df": data})
+                # Using safe_eval instead of eval for security
+                rule_matched = self.safe_eval(rule.condition, {"df": data})
                 results[rule_name] = {
                     "rule_matched": rule_matched,
                     "action": rule.action if rule_matched else None,
@@ -266,6 +280,26 @@ class ClinicalRulesEngine:
                 }
         return results
 
+    @staticmethod
+    def safe_eval(expr, variables):
+        """
+        Safely evaluate an expression using AST parsing.
+        Only allows certain node types to prevent execution of arbitrary code.
+        """
+        allowed_nodes = (
+            ast.Expression, ast.BoolOp, ast.BinOp, ast.UnaryOp, ast.Compare,
+            ast.Call, ast.Name, ast.Load, ast.Constant, ast.Num, ast.Str,
+            ast.List, ast.Tuple, ast.Dict
+        )
+        try:
+            node = ast.parse(expr, mode='eval')
+            for subnode in ast.walk(node):
+                if not isinstance(subnode, allowed_nodes):
+                    raise ValueError(f"Unsupported expression: {expr}")
+            return eval(compile(node, '<string>', mode='eval'), {"__builtins__": None}, variables)
+        except Exception as e:
+            raise ValueError(f"Invalid expression: {e}")
+
 class ClinicalKPI(BaseModel):
     """Define a clinical KPI."""
     name: str
@@ -284,8 +318,8 @@ class ClinicalKPIMonitoring:
         results = {}
         for kpi_name, kpi in self.kpis.items():
             try:
-                # **Warning**: Using eval can be dangerous. Ensure that user inputs are sanitized.
-                kpi_value = eval(kpi.calculation, {"__builtins__": None}, {"df": data})
+                # Using safe_eval instead of eval for security
+                kpi_value = self.safe_eval(kpi.calculation, {"df": data})
                 status = self.evaluate_threshold(kpi_value, kpi.threshold)
                 results[kpi_name] = {
                     "value": kpi_value,
@@ -305,6 +339,26 @@ class ClinicalKPIMonitoring:
         except TypeError:
             return "Threshold Evaluation Not Applicable"
 
+    @staticmethod
+    def safe_eval(expr, variables):
+        """
+        Safely evaluate an expression using AST parsing.
+        Only allows certain node types to prevent execution of arbitrary code.
+        """
+        allowed_nodes = (
+            ast.Expression, ast.BoolOp, ast.BinOp, ast.UnaryOp, ast.Compare,
+            ast.Call, ast.Name, ast.Load, ast.Constant, ast.Num, ast.Str,
+            ast.List, ast.Tuple, ast.Dict
+        )
+        try:
+            node = ast.parse(expr, mode='eval')
+            for subnode in ast.walk(node):
+                if not isinstance(subnode, allowed_nodes):
+                    raise ValueError(f"Unsupported expression: {expr}")
+            return eval(compile(node, '<string>', mode='eval'), {"__builtins__": None}, variables)
+        except Exception as e:
+            raise ValueError(f"Invalid expression: {e}")
+
 class DiagnosisSupport(ABC):
     """Abstract class for implementing clinical diagnoses."""
     @abstractmethod
@@ -394,21 +448,67 @@ class MedicalKnowledgeBase(ABC):
         pass
 
 class SimpleMedicalKnowledge(MedicalKnowledgeBase):
-    """Simple Medical Knowledge Class with TF-IDF, NER, and PubMed."""
+    """Enhanced Medical Knowledge Class using Groq API."""
     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.",
-            "cancer drugs": "Please consult with your doctor to get personalized treatment and the latest drug information for your type of cancer."
+        self.api_endpoint = GROQ_API_ENDPOINT
+        self.api_key = GROQ_API_KEY
+        self.pub_email = PUB_EMAIL
+
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
         }
-        self.vectorizer = TfidfVectorizer()
-        self.tfidf_matrix = self.vectorizer.fit_transform(self.knowledge_base.values())
+        
+        # Initialize spaCy model for entity recognition if needed
+        try:
+            self.nlp = spacy.load("en_core_web_sm")
+        except OSError:
+            import subprocess
+            import sys
+            subprocess.run([sys.executable, "-m", "spacy", "download", "en_core_web_sm"])
+            self.nlp = spacy.load("en_core_web_sm")
+
+    def search_medical_info(self, query: str, pub_email: str = "") -> str:
+        """
+        Uses the Groq API to fetch medical information based on the user's query.
+        """
+        try:
+            # Preprocess the query if necessary (e.g., entity recognition)
+            doc = self.nlp(query.lower())
+            entities = [ent.text for ent in doc.ents]
+            if entities:
+                processed_query = " ".join(entities)
+            else:
+                processed_query = query.lower()
+
+            # Prepare the payload for the Groq API
+            payload = {
+                "query": processed_query,
+                "context": "medical"  # Assuming the API can handle context specification
+            }
+
+            # Make the API request
+            response = requests.post(
+                self.api_endpoint,
+                headers=self.headers,
+                data=json.dumps(payload)
+            )
 
-    def search_pubmed(self, query: str, email: str) -> str:
-        """Search PubMed for abstracts related to the query."""
+            if response.status_code == 200:
+                data = response.json()
+                answer = data.get("answer", "I'm sorry, I couldn't find relevant information.")
+                pubmed_abstract = self.fetch_pubmed_abstract(processed_query, pub_email)
+                return f"**Based on your query:** {answer}\n\n**PubMed Abstract:**\n\n{pubmed_abstract}"
+            else:
+                return f"Error: Received status code {response.status_code} from Groq API."
+
+        except Exception as e:
+            return f"Medical Knowledge Search Failed: {str(e)}"
+
+    def fetch_pubmed_abstract(self, query: str, email: str) -> str:
+        """
+        Searches PubMed for abstracts related to the query.
+        """
         try:
             Entrez.email = email
             handle = Entrez.esearch(db="pubmed", term=query, retmax=1, sort='relevance')
@@ -424,68 +524,6 @@ class SimpleMedicalKnowledge(MedicalKnowledgeBase):
         except Exception as e:
             return f"Error searching PubMed: {e}"
 
-    def search_medical_info(self, query: str, pub_email: str = "") -> str:
-        """Search the medical knowledge base and PubMed for relevant information."""
-        try:
-            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:
-                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.")
-
-            # Enhanced PubMed Search: Combine query and best_match_keyword for better relevance
-            pubmed_query = f"{query_lower} AND {best_match_keyword}"
-            pubmed_result = self.search_pubmed(pubmed_query, pub_email)
-
-            feedback_key = f"feedback_{query_lower}"  # Creating a unique key for feedback
-
-            response = f"**Based on your query:** {best_match_info}\n\n"
-
-            if "Error searching PubMed" not in pubmed_result and "No abstracts found" not in pubmed_result:
-                # Format the PubMed abstract with proper markdown
-                abstract_title = pubmed_result.split('\n')[0]  # Assuming the first line is the title
-                abstract_body = '\n'.join(pubmed_result.split('\n')[2:])  # Skipping authors and affiliations
-                response += f"**PubMed Abstract:**\n\n**{abstract_title}**\n\n{abstract_body}"
-            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 searching PubMed" 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):
@@ -711,7 +749,7 @@ def initialize_session_state():
     if 'knowledge_base' not in st.session_state:
         st.session_state.knowledge_base = SimpleMedicalKnowledge()
     if 'pub_email' not in st.session_state:
-        st.session_state.pub_email = st.secrets.get("PUB_EMAIL", "")  # Load PUB_EMAIL from secrets
+        st.session_state.pub_email = PUB_EMAIL  # Load PUB_EMAIL from environment variables
 
 def data_management_section():
     """Handles the data management section in the sidebar."""