Update app.py

app.py

@@ -1,19 +1,79 @@
 import gradio as gr
-import os
-from components import pubmed_search
-from components import model_utils
-import time
+from Bio import Entrez
+import os  # For environment variables and file paths
+from components import federated_learning
+
 # ---------------------------- Configuration ----------------------------
-ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "ENTREZ_EMAIL")
-HUGGINGFACE_API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN", "HUGGINGFACE_API_TOKEN")
+ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "ENTREZ_EMAIL")  # Use environment variable, default fallback
+HUGGINGFACE_API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN", "HUGGINGFACE_API_TOKEN") # Use environment variable, default fallback
+
+# ---------------------------- Helper Functions ----------------------------
+
+def log_error(message: str):
+    """Logs an error message to the console and a file (if possible)."""
+    print(f"ERROR: {message}")
+    try:
+        with open("error_log.txt", "a") as f:
+            f.write(f"{message}\n")
+    except:
+        print("Couldn't write to error log file.")  #If logging fails, still print to console
+
+# ---------------------------- Tool Functions ----------------------------
+
+def search_pubmed(query: str) -> list:
+    """Searches PubMed and returns a list of article IDs."""
+    try:
+        Entrez.email = ENTREZ_EMAIL
+        handle = Entrez.esearch(db="pubmed", term=query, retmax="5")
+        record = Entrez.read(handle)
+        handle.close()
+        return record["IdList"]
+    except Exception as e:
+        log_error(f"PubMed search error: {e}")
+        return [f"Error during PubMed search: {e}"]
+
+def fetch_abstract(article_id: str) -> str:
+    """Fetches the abstract for a given PubMed article ID."""
+    try:
+        Entrez.email = ENTREZ_EMAIL
+        handle = Entrez.efetch(db="pubmed", id=article_id, rettype="abstract", retmode="text")
+        abstract = handle.read()
+        handle.close()
+        return abstract
+    except Exception as e:
+        log_error(f"Error fetching abstract for {article_id}: {e}")
+        return f"Error fetching abstract for {article_id}: {e}"
+
+# ---------------------------- Agent Function ----------------------------
+
+def medai_agent(query: str) -> str:
+    """Orchestrates the medical literature review and presents abstract."""
+    article_ids = search_pubmed(query)
 
-# ---------------------------- Global Variables ----------------------------
+    if isinstance(article_ids, list) and article_ids:
+        results = []
+        for article_id in article_ids:
+            abstract = fetch_abstract(article_id)
+            if "Error" not in abstract:
+                results.append(f"<div class='article'>\n"
+                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
+                               f"  <p class='abstract'><strong>Abstract:</strong> {abstract}</p>\n"
+                               f"</div>\n")
+            else:
+                results.append(f"<div class='article error'>\n"
+                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
+                               f"  <p class='error-message'>Error processing article: {abstract}</p>\n"
+                               f"</div>\n")
+        return "\n".join(results)
+    else:
+        return f"No articles found or error occurred: {article_ids}"
 
 # ---------------------------- Gradio Interface ----------------------------
 
 def launch_gradio():
     """Launches the Gradio interface."""
 
+    # CSS to style the article output
     css = """
     .article {
       border: 1px solid #ddd;
@@ -40,12 +100,15 @@ def launch_gradio():
         gr.Markdown("# MedAI: Medical Literature Review")
         gr.Markdown("Enter a medical query to retrieve abstracts from PubMed.")
 
-        query_input = gr.Textbox(lines=3, placeholder="Enter your medical query (e.g., 'new treatments for diabetes')...")
+        query_input = gr.Textbox(lines=3, placeholder="Enter your medical query to get abstract from PubMed.")
         submit_button = gr.Button("Submit")
         output_results = gr.HTML()  # Use HTML for formatted output
+        federated_learning_output = gr.HTML()
 
         # Get data
-        submit_button.click(pubmed_search.medai_agent, inputs=query_input, outputs=output_results)
+        submit_button.click(medai_agent, inputs=query_input, outputs=output_results)
+        run_fl_button = gr.Button("Run Federated Learning (Conceptual)")
+        run_fl_button.click(federated_learning.run_federated_learning, outputs = federated_learning_output)
 
     iface.launch()