Update retrieval.py

retrieval.py

@@ -7,7 +7,7 @@ import chromadb
 from chromadb.config import Settings
 from transformers import AutoTokenizer, AutoModel
 
-# Optional: Set your PubMed API key from environment variables
+# Optional: Set your PubMed API key from environment variables.
 PUBMED_API_KEY = os.environ.get("PUBMED_API_KEY", "<YOUR_NCBI_API_KEY>")
 
 #############################################
@@ -15,8 +15,8 @@ PUBMED_API_KEY = os.environ.get("PUBMED_API_KEY", "<YOUR_NCBI_API_KEY>")
 #############################################
 def fetch_pubmed_abstracts(query: str, max_results: int = 5) -> List[str]:
     """
-    Fetches PubMed abstracts for the given query using NCBI's E-utilities.
-    Returns a list of abstract texts.
+    Retrieves PubMed abstracts for a given clinical query using NCBI's E-utilities.
+    Designed to quickly fetch up to 'max_results' abstracts.
     """
     search_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
     params = {
@@ -26,7 +26,7 @@ def fetch_pubmed_abstracts(query: str, max_results: int = 5) -> List[str]:
         "api_key": PUBMED_API_KEY,
         "retmode": "json"
     }
-    r = requests.get(search_url, params=params)
+    r = requests.get(search_url, params=params, timeout=10)
     r.raise_for_status()
     data = r.json()
 
@@ -41,7 +41,7 @@ def fetch_pubmed_abstracts(query: str, max_results: int = 5) -> List[str]:
             "retmode": "text",
             "api_key": PUBMED_API_KEY
         }
-        fetch_resp = requests.get(fetch_url, params=fetch_params)
+        fetch_resp = requests.get(fetch_url, params=fetch_params, timeout=10)
         fetch_resp.raise_for_status()
         abstract_text = fetch_resp.text.strip()
         if abstract_text:
@@ -53,7 +53,8 @@ def fetch_pubmed_abstracts(query: str, max_results: int = 5) -> List[str]:
 #############################################
 class EmbedFunction:
     """
-    Wraps a Hugging Face embedding model to produce embeddings for a list of strings.
+    Uses a Hugging Face embedding model to generate embeddings for a list of strings.
+    This function is crucial for indexing abstracts for similarity search.
     """
     def __init__(self, model_name: str):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -80,7 +81,7 @@ class EmbedFunction:
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 embed_function = EmbedFunction(EMBED_MODEL_NAME)
 
-# Use a temporary directory for persistent storage.
+# Use a temporary directory for persistent storage to ensure a fresh initialization.
 temp_dir = tempfile.mkdtemp()
 print("Using temporary persist_directory:", temp_dir)
 
@@ -91,24 +92,24 @@ client = chromadb.Client(
     )
 )
 
-# Create or get the collection. Use a clear name.
+# Create or retrieve the collection for medical abstracts.
 collection = client.get_or_create_collection(
     name="ai_medical_knowledge",
     embedding_function=embed_function
 )
 
-# Force initialization: add a dummy document and perform a dummy query.
+# Optional: Force initialization with a dummy document to ensure the schema is set up.
 try:
     collection.add(documents=["dummy"], ids=["dummy"])
     _ = collection.query(query_texts=["dummy"], n_results=1)
-    # Optionally, remove the dummy document if needed (Chromadb might not support deletion, so you can ignore it)
     print("Dummy initialization successful.")
 except Exception as init_err:
     print("Dummy initialization failed:", init_err)
 
 def index_pubmed_docs(docs: List[str], prefix: str = "doc"):
     """
-    Adds documents to the Chromadb collection with unique IDs.
+    Indexes PubMed abstracts into the Chroma vector store.
+    Each document is assigned a unique ID based on the query prefix.
     """
     for i, doc in enumerate(docs):
         if doc.strip():
@@ -122,7 +123,8 @@ def index_pubmed_docs(docs: List[str], prefix: str = "doc"):
 
 def query_similar_docs(query: str, top_k: int = 3) -> List[str]:
     """
-    Retrieves the top_k similar documents from Chromadb based on embedding similarity.
+    Searches the indexed abstracts for those most similar to the given query.
+    Returns the top 'top_k' documents.
     """
     results = collection.query(query_texts=[query], n_results=top_k)
     return results["documents"][0] if results and results["documents"] else []
@@ -132,10 +134,12 @@ def query_similar_docs(query: str, top_k: int = 3) -> List[str]:
 #############################################
 def get_relevant_pubmed_docs(user_query: str) -> List[str]:
     """
-    End-to-end pipeline:
+    Complete retrieval pipeline:
       1. Fetch PubMed abstracts for the query.
-      2. Index them in Chromadb.
-      3. Retrieve the top relevant documents.
+      2. Index the abstracts into the vector store.
+      3. Retrieve and return the most similar documents.
+    
+    Designed for clinicians to quickly access relevant literature.
     """
     new_abstracts = fetch_pubmed_abstracts(user_query, max_results=5)
     if not new_abstracts: