Update pubmed_rag.py

pubmed_rag.py

@@ -23,7 +23,7 @@ def search_pubmed(query, max_results=MAX_PUBMED_RESULTS):
         "tool": "AdvancedMedicalAI",
         "email": MY_PUBMED_EMAIL,
     }
-    response = requests.get(url, params=params)
+    response = requests.get(url, params=params, timeout=10)
     response.raise_for_status()
     data = response.json()
     return data.get("esearchresult", {}).get("idlist", [])
@@ -41,7 +41,7 @@ def fetch_abstract(pmid):
         "tool": "AdvancedMedicalAI",
         "email": MY_PUBMED_EMAIL,
     }
-    response = requests.get(url, params=params)
+    response = requests.get(url, params=params, timeout=10)
     response.raise_for_status()
     return response.text.strip()
 
@@ -52,7 +52,8 @@ def fetch_pubmed_abstracts(pmids):
     results = {}
     for pmid in pmids:
         try:
-            results[pmid] = fetch_abstract(pmid)
+            abstract = fetch_abstract(pmid)
+            results[pmid] = abstract
         except Exception as e:
             results[pmid] = f"Error fetching PMID {pmid}: {e}"
     return results
@@ -78,8 +79,8 @@ def summarize_text(text, chunk_size=SUMMARIZATION_CHUNK_SIZE):
     if current_chunk:
         chunks.append(" ".join(current_chunk))
 
-    summaries = [
-        summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-        for chunk in chunks
-    ]
-    return " ".join(summaries)
+    summaries = []
+    for chunk in chunks:
+        summary = summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
+        summaries.append(summary)
+    return " ".join(summaries)