Update app.py

app.py

@@ -18,10 +18,14 @@ import os
 import numpy as np
 from scipy.stats import ttest_ind, f_oneway
 import json
+from Bio import Entrez
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 
 # Initialize Groq Client
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
+
 # ---------------------- Base Classes and Schemas ---------------------------
 class ResearchInput(BaseModel):
     """Base schema for research tool inputs"""
@@ -287,17 +291,50 @@ class MedicalKnowledgeBase():
     pass
 
 class SimpleMedicalKnowledge(MedicalKnowledgeBase):
-    """Simple Medical Knowledge Class"""
-    def search_medical_info(self, query: str) -> str:
+    """Simple Medical Knowledge Class with TF-IDF 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."
+       }
+       self.vectorizer = TfidfVectorizer()
+       self.tfidf_matrix = self.vectorizer.fit_transform(self.knowledge_base.values())
+       
+    def search_pubmed(self, query: str, email: str) -> str:
+      try:
+          Entrez.email = email
+          handle = Entrez.esearch(db="pubmed", term=query, retmax=1)
+          record = Entrez.read(handle)
+          handle.close()
+          if record["IdList"]:
+            handle = Entrez.efetch(db="pubmed", id=record["IdList"][0], rettype="abstract", retmode="text")
+            abstract = handle.read()
+            handle.close()
+            return abstract
+          else:
+            return "No abstracts found for this query on PubMed"
+      except Exception as e:
+            return f"Error searching pubmed {e}"
+
+
+    def search_medical_info(self, query: str, pub_email: 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"
+          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 for this query on PubMed" not in pubmed_result:
+             return f"Based on the query provided, I found this: {best_match_info} \n\nFrom Pubmed I also found the following abstract: \n {pubmed_result}"
           else:
-              return "No specific information is available"
+            return f"Based on the query provided, I found this: {best_match_info} \n\n{pubmed_result}"
         except Exception as e:
-          return f"Medical Knowledge Search Failed {e}"
+            return f"Medical Knowledge Search Failed {e}"
 
 
 class ForecastingEngine(ABC):
@@ -494,8 +531,9 @@ def main():
        st.session_state.treatment_recommendation = BasicTreatmentRecommendation()
     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
   
-
     # Sidebar for Data Management
     with st.sidebar:
         st.header("⚙️ Data Management")
@@ -668,51 +706,4 @@ 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, 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()
+