Added rag features

Author: Shyaam-04

backend/chroma_db/a26d24c1-363d-4e7e-927b-5594ca985819/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:587f223a00318da2e6181e47be8e8d133cb354c6d10a3bf2b61a0eec70bea177
+size 22484000

backend/chroma_db/a26d24c1-363d-4e7e-927b-5594ca985819/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a545742ed5f9e3ad7971a65470e401ab943cb853438d631a10135760fc124587
+size 100

backend/chroma_db/a26d24c1-363d-4e7e-927b-5594ca985819/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb6f636174c8836330257a4107c0b4d5276da7ef60569063ad49849cccd0074c
+size 220694

backend/chroma_db/a26d24c1-363d-4e7e-927b-5594ca985819/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcfdc8595cfc4464d9c91ebba31facc865e9014e9a97ebc9c574d58b758dacb9
+size 28000

backend/chroma_db/a26d24c1-363d-4e7e-927b-5594ca985819/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59552bb1d187ffe10aa708706e68b72eb6e28f1916034dde682211a8a0cb4c63
+size 61524

backend/chroma_db/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b06cc3f72e27da6221f33d972c76f898b9d77b91148689ee27ce9603f72ade36
+size 121520128

backend/embedder.py

@@ -0,0 +1,12 @@
+from sentence_transformers import SentenceTransformer
+
+# Load once at module level — not on every call
+model = SentenceTransformer("pritamdeka/S-PubMedBert-MS-MARCO")
+
+def embed(text: str) -> list[float]:
+    """Convert text to embedding vector."""
+    return model.encode(text, convert_to_numpy=True).tolist()
+
+def embed_batch(texts: list[str]) -> list[list[float]]:
+    """Convert a batch of texts to embedding vectors."""
+    return model.encode(texts, convert_to_numpy=True, batch_size=32).tolist()

backend/engine.py

@@ -1,6 +1,7 @@
 from itertools import combinations
 import requests
 import json
+from retriever import retrieve
 
 DRUG_CLASS_MAP = {
     "penicillin": ["amoxicillin", "ampicillin", "penicillin", "flucloxacillin"],
@@ -12,85 +13,88 @@
 
 
 def load_fallback_data():
-    with open("data/fallback_interactions.json","r") as f:
+    with open("data/fallback_interactions.json", "r") as f:
         data = json.load(f)
     return data
 
+
 def check_interactions_fallback(medicines: list[str], fallback_data: list) -> list:
-    pairs = list(combinations(medicines,2))
+    pairs = list(combinations(medicines, 2))
     results = []
     for pair in pairs:
         drug_1 = pair[0]
         drug_2 = pair[1]
-
         for entry in fallback_data:
-            if (entry["drug_a"].lower() == drug_1.lower() and entry["drug_b"].lower() == drug_2.lower()) or (entry["drug_a"].lower() == drug_2.lower() and entry["drug_b"].lower() == drug_1.lower()):
+            if (
+                entry["drug_a"].lower() == drug_1.lower()
+                and entry["drug_b"].lower() == drug_2.lower()
+            ) or (
+                entry["drug_a"].lower() == drug_2.lower()
+                and entry["drug_b"].lower() == drug_1.lower()
+            ):
                 results.append(entry)
-
     return results
 
+
 def check_allergies(medicines: list[str], known_allergies: list[str]) -> list:
-    # For each medicine, check if it matches any known allergy
-    # Return a list of allergy alerts
     matches = []
     for medicine in medicines:
         for known_allergy in known_allergies:
-            #exact matching
             if medicine.lower() == known_allergy.lower():
                 matches.append({
                     "medicine": medicine,
                     "reason": f"Patient is allergic to {known_allergy}",
                     "severity": "high"
                 })
-
-            #mapping with drug classes
             else:
                 for drug_class, drug_items in DRUG_CLASS_MAP.items():
-                    if known_allergy.lower() == drug_class.lower() and medicine.lower() in drug_items:
+                    if (
+                        known_allergy.lower() == drug_class.lower()
+                        and medicine.lower() in drug_items
+                    ):
                         matches.append({
                             "medicine": medicine,
                             "reason": f"{drug_class} class",
                             "severity": "high"
                         })
-
     return matches
 
+
 def parse_llm_response(content: str) -> dict:
     try:
         content = content.strip()
-        #removing md output if any by the llm
         if content.startswith("```"):
             content = content.split("```")[1]
             if content.startswith("json"):
                 content = content[4:]
 
-        #finding { } brackets in json to extract pure json
-
         start = content.index("{")
-        end = content.rindex("}")+1
+        end = content.rindex("}") + 1
         content = content[start:end]
 
-        #loading content
         data = json.loads(content)
-
-        #ensuring required fields atleast contain default values
-        data.setdefault("interactions",[])
-        data.setdefault("allergy_alerts",[])
-        data.setdefault("requires_doctor_review",True)
+        data.setdefault("interactions", [])
+        data.setdefault("allergy_alerts", [])
+        data.setdefault("requires_doctor_review", True)
 
         return data
-    
+
     except Exception:
         return {
             "interactions": [],
             "allergy_alerts": [],
             "requires_doctor_review": True
         }
-    
+
+
 def check_interactions_llm(medicines: list[str], patient_history) -> dict:
     with open("prompts/system_prompt.txt", "r") as f:
         system_prompt = f.read()
 
+    # Retrieve relevant clinical context from knowledge base
+    retrieved_context = retrieve(medicines, top_k=10)
+
+    # Build user message with retrieved evidence
     user_message = json.dumps({
         "medicines": medicines,
         "patient_history": {
@@ -99,24 +103,53 @@ def check_interactions_llm(medicines: list[str], patient_history) -> dict:
             "known_allergies": patient_history.known_allergies,
             "current_medications": patient_history.current_medications,
             "conditions": patient_history.conditions
-        }
+        },
+        "clinical_reference_data": retrieved_context,
+        "instruction": "Analyze ONLY the medicines listed above. Use clinical_reference_data as supporting evidence only. Return ONLY valid JSON matching the required output schema. Do not summarize the reference data."
     }, indent=2)
 
+
     response = requests.post(
         "http://localhost:11434/api/chat",
         json={
             "model": "qwen2.5",
             "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_message}
+                {
+                    "role": "system",
+                    "content": system_prompt
+                },
+                {
+                    "role": "user",
+                    "content": f"CLINICAL REFERENCE DATA (use as evidence only, do not summarize):\n\n{retrieved_context}"
+                },
+                {
+                    "role": "assistant",
+                    "content": "I have reviewed the clinical reference data. I will use it as supporting evidence only. I am ready to analyze the medicines you specify and return JSON in the required schema."
+                },
+                {
+                    "role": "user",
+                    "content": f"Now analyze this prescription and return ONLY valid JSON:\n\n{json.dumps({'medicines': medicines, 'patient_history': {'age': patient_history.age, 'weight': patient_history.weight, 'known_allergies': patient_history.known_allergies, 'current_medications': patient_history.current_medications, 'conditions': patient_history.conditions}}, indent=2)}"
+                }
             ],
             "stream": False
         },
-        timeout = 30
+        timeout=120
     )
 
     response.raise_for_status()
-
     content = response.json()["message"]["content"]
     return parse_llm_response(content)
 
+
+if __name__ == "__main__":
+    class MockHistory:
+        age = "80"
+        weight = "60"
+        known_allergies = []
+        current_medications = []
+        conditions = ["Atrial Fibrillation"]
+
+    medicines = ["Digoxin","Verapamil"]
+    result = check_interactions_llm(medicines, MockHistory())
+    print("FINAL RESULT:")
+    print(json.dumps(result, indent=2))

backend/ingest.py

@@ -0,0 +1,125 @@
+import json
+import os
+import pickle
+from rank_bm25 import BM25Okapi
+import chromadb
+from embedder import embed_batch
+
+# ── Paths ──────────────────────────────────────────────
+DATA_PATH = "data/drug-label-0001-of-0013.json"
+CHROMA_PATH = "chroma_db"
+BM25_PATH = "bm25_index.pkl"
+
+# ── ChromaDB setup ─────────────────────────────────────
+client = chromadb.PersistentClient(path=CHROMA_PATH)
+collection = client.get_or_create_collection(
+    name="drug_interactions",
+    metadata={"hnsw:space": "cosine"}
+)
+
+def load_records(path: str) -> list[dict]:
+    """Load only records that have drug_interactions field."""
+    print("Loading bulk data...")
+    with open(path, "r") as f:
+        data = json.load(f)
+
+    records = [
+        r for r in data["results"]
+        if "drug_interactions" in r and r["drug_interactions"]
+    ]
+    print(f"Found {len(records)} records with interaction data")
+    return records
+
+def extract_drug_name(record: dict) -> str:
+    openfda = record.get("openfda", {})
+    generic = openfda.get("generic_name", [])
+    brand = openfda.get("brand_name", [])
+    spl = record.get("spl_product_data_elements", [])
+
+    if generic:
+        return generic[0].upper()
+    elif brand:
+        return brand[0].upper()
+    elif spl:
+        return spl[0].split()[0].upper()
+    else:
+        return "UNKNOWN"
+
+def chunk_text(text: str, drug_name: str, chunk_size: int = 600, overlap: int = 100) -> list[str]:
+    """Split interaction text into overlapping chunks."""
+    words = text.split()
+    chunks = []
+    start = 0
+
+    while start < len(words):
+        end = start + chunk_size
+        chunk = " ".join(words[start:end])
+        # Prepend drug name to every chunk for BM25 keyword matching
+        chunks.append(f"{drug_name}: {chunk}")
+        start += chunk_size - overlap
+
+    return chunks
+
+def ingest():
+    # ── Check if already ingested ──────────────────────
+    existing = collection.count()
+    if existing > 0:
+        print(f"ChromaDB already has {existing} chunks. Skipping ingestion.")
+        print("Delete chroma_db/ folder to re-ingest.")
+        return
+
+    records = load_records(DATA_PATH)
+    
+    all_chunks = []
+    all_ids = []
+    all_metadata = []
+
+    print("Chunking records...")
+    for i, record in enumerate(records):
+        drug_name = extract_drug_name(record)
+        interaction_text = record["drug_interactions"][0]
+        chunks = chunk_text(interaction_text, drug_name)
+
+        for j, chunk in enumerate(chunks):
+            all_chunks.append(chunk)
+            all_ids.append(f"{i}_{j}")
+            all_metadata.append({"drug_name": drug_name})
+
+    print(f"Total chunks: {len(all_chunks)}")
+
+    # ── Embed and store in ChromaDB in batches ─────────
+    print("Embedding and storing in ChromaDB...")
+    batch_size = 50
+
+    for start in range(0, len(all_chunks), batch_size):
+        end = start + batch_size
+        batch_chunks = all_chunks[start:end]
+        batch_ids = all_ids[start:end]
+        batch_metadata = all_metadata[start:end]
+
+        embeddings = embed_batch(batch_chunks)
+
+        collection.add(
+            documents=batch_chunks,
+            embeddings=embeddings,
+            ids=batch_ids,
+            metadatas=batch_metadata
+        )
+
+        if start % 500 == 0:
+            print(f"  Processed {start}/{len(all_chunks)} chunks...")
+
+    # ── Build and save BM25 index ──────────────────────
+    print("Building BM25 index...")
+    tokenized = [chunk.lower().split() for chunk in all_chunks]
+    bm25 = BM25Okapi(tokenized)
+
+    with open(BM25_PATH, "wb") as f:
+        pickle.dump((bm25, all_chunks), f)
+
+    print(f"\nIngestion complete.")
+    print(f"ChromaDB chunks: {collection.count()}")
+    print(f"BM25 index saved to {BM25_PATH}")
+
+if __name__ == "__main__":
+    ingest()