Upload 6 files

Dockerfile

@@ -0,0 +1,28 @@
+FROM python:3.10.13-slim
+
+WORKDIR /app
+
+# System dependencies required for blis / thinc / spacy
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    gcc \
+    g++ \
+    git \
+    libopenblas-dev \
+    libomp-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Upgrade pip first
+RUN pip install --upgrade pip setuptools wheel
+
+# Install Python deps
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy app
+COPY . .
+
+# Streamlit default port
+EXPOSE 7860
+
+CMD ["streamlit", "run", "app.py", "--server.port=7860", "--server.address=0.0.0.0"]

app.py

@@ -0,0 +1,213 @@
+import sys
+import os
+import spacy
+from groq import Groq
+from dotenv import load_dotenv
+
+load_dotenv()  # This loads variables from .env into environment
+
+dotenv_path = os.path.join(os.path.dirname(__file__), 'API_key.env')  # adjust if needed
+load_dotenv(dotenv_path)
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'pyspur/backend/')))
+
+from embedding import discharge_collection, trials_collection, get_embedding
+
+from serpapi import GoogleSearch
+from pyspur.backend.pyspur.nodes.decorator import tool_function
+
+# Load API key from Hugging Face secret
+
+groq_api_key = os.getenv("GROQ_API_KEY")
+if not groq_api_key:
+    raise ValueError("Missing GROQ_API_KEY in environment variables.")
+
+serp_api_key = os.environ.get("SERP_API_KEY")
+if not serp_api_key:
+    raise ValueError("Missing SERP_API_KEY in environment variables.")
+
+# Initialize LLM client and spaCy
+
+client = Groq(api_key=groq_api_key)
+
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    from spacy.cli import download
+    download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
+
+SYMPTOM_HINTS = [
+    "chest pain", "shortness of breath", "fatigue", "dizziness",
+    "nausea", "vomiting", "palpitations", "sweating", "jaw pain",
+    "arm pain", "back pain", "tightness", "pressure in chest",
+    "arrhythmia", "tachycardia", "bradycardia", "angina",
+    "edema", "dyspnea", "syncope", "lightheadedness",
+    "ejection fraction", "myocardial infarction", "heart failure",
+    "cardiomyopathy", "cardiac arrest"
+]
+
+@tool_function(name="chat_memory_tool")
+def chat_memory_tool(memory: str, model: str = "llama-3.3-70b-versatile") -> str:
+    doc = nlp(memory)
+    found_symptoms = set(
+        keyword for chunk in doc.noun_chunks for keyword in SYMPTOM_HINTS if keyword in chunk.text.lower()
+    )
+    symptom_context = (
+        f"Previously mentioned symptoms include: {', '.join(found_symptoms)}."
+        if found_symptoms else "No clear symptoms found in memory."
+    )
+    response = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": "You are a medical assistant summarizing prior symptoms from memory."},
+            {"role": "assistant", "content": memory},
+            {"role": "user", "content": (
+                f"The patient previously reported: {memory}\n\n"
+                f"Symptoms extracted: {symptom_context}\n"
+                "Please provide a clear, concise, and helpful summary of these symptoms and suggest next steps."
+            )}
+        ]
+    )
+    return response.choices[0].message.content
+
+@tool_function(name="treatment_tool")
+def treatment_tool(query: str, model: str = "llama-3.3-70b-versatile", use_rag: bool = True) -> str:
+    try:
+        query_embedding = get_embedding(query)
+        if use_rag:
+            results = discharge_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=5,
+                include=["documents"]
+            )
+            top_docs = results['documents'][0] if results and results['documents'] else []
+            top_docs = [doc[:1500] for doc in top_docs]
+            combined_context = "\n\n".join(top_docs)
+            prompt = (
+                "You are a helpful medical assistant. Based on the following discharge notes, "
+                "recommend essential treatment.\n\n"
+                f"### Notes:\n{combined_context}\n\n### Condition:\n{query}"
+            )
+        else:
+            prompt = f"Patient condition: {query}. What treatment is recommended?"
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a medically accurate and safety-focused clinical assistant."},
+                {"role": "user", "content": prompt}
+            ]
+        )
+        return response.choices[0].message.content
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool_function(name="symptom_search_tool")
+def symptom_search_tool(symptom_description: str, model: str = "llama-3.3-70b-versatile") -> str:
+    def perform_search(query):
+        params = {
+            "engine": "google",
+            "q": f"{query} possible causes site:mayoclinic.org OR site:webmd.com OR site:nih.gov",
+            "api_key": serp_api_key
+        }
+        return GoogleSearch(params).get_dict().get("organic_results", [])
+
+    try:
+        results = perform_search(symptom_description)
+        if not results:
+            return "No reliable medical source found."
+
+        sources = []
+        snippets_with_citations = []
+        for res in results[:3]:
+            if 'snippet' in res and 'link' in res:
+                source_url = res['link']
+                domain = source_url.split("//")[-1].split("/")[0].replace("www.", "")
+                snippets_with_citations.append(f"{res['snippet']} (Source: {domain})")
+                sources.append(source_url)
+
+        search_context = "\n\n".join(snippets_with_citations)
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a medical assistant using trusted web sources to explain symptom causes."},
+                {"role": "assistant", "content": search_context},
+                {"role": "user", "content": f"What could be the cause of: {symptom_description}?"}
+            ]
+        )
+
+        bulleted_sources = "\n".join(f"- {url}" for url in sources)
+        return response.choices[0].message.content + "\n\n**Sources:**\n" + bulleted_sources
+
+    except Exception as e:
+        return f"Search error: {str(e)}"
+
+@tool_function(name="trial_matcher_tool")
+def trial_matcher_tool(discharge_note: str, model: str = "llama-3.3-70b-versatile", use_rag: bool = True) -> str:
+    try:
+        query_embedding = get_embedding(discharge_note)
+        results = trials_collection.query(
+            query_embeddings=[query_embedding],
+            n_results=3,
+            include=["documents", "metadatas"]
+        )
+        if not results.get('documents') or not results['documents'][0]:
+            return "No matching clinical trials were found for the provided note."
+
+        summaries = []
+        for i, (doc, meta) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
+            nct_id = meta.get("NCT ID") or "Unknown ID"
+            truncated_doc = doc.strip()[:1500]
+            if use_rag:
+                summary_prompt = (
+                    f"You are a clinical assistant reviewing a matched clinical trial.\n"
+                    f"Summarize the trial using **bullet points only** for the following fields:\n"
+                    f"- NCT ID\n- Study Title\n- Conditions\n- Inclusion Criteria\n- Exclusion Criteria\n\n"
+                    f"Use bullets under each field. Maintain a clean format. Respond only with the summary.\n\n"
+                    f"Trial Description:\nNCT ID: {nct_id}\n{truncated_doc}"
+                )
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=[
+                        {"role": "system", "content": "You are a medically precise clinical research assistant."},
+                        {"role": "user", "content": summary_prompt}
+                    ]
+                )
+                summaries.append(f"### Trial {i+1}:\n{response.choices[0].message.content}")
+            else:
+                summaries.append(f"### Trial {i+1}:\nNCT ID: {nct_id}\n\n{truncated_doc}")
+
+        return "\n\n---\n\n".join(summaries)
+
+    except Exception as e:
+        return f"Error during trial matching: {str(e)}"
+
+# Tool routing via keyword logic
+
+TOOL_ROUTER = {
+    "symptom": ("symptom_search_tool", False),
+    "treatment": ("treatment_tool", True),
+    "trial": ("trial_matcher_tool", True)
+}
+
+TOOL_FUNCTIONS = {
+    "chat_memory_tool": chat_memory_tool,
+    "treatment_tool": treatment_tool,
+    "symptom_search_tool": symptom_search_tool,
+    "trial_matcher_tool": trial_matcher_tool
+}
+
+def run_tool(query: str, model: str, use_rag: bool) -> str:
+    for keyword, (tool_name, supports_rag) in TOOL_ROUTER.items():
+        if keyword in query.lower():
+            print(f"Tool selected by PySpur: {tool_name}")
+            tool_func = TOOL_FUNCTIONS[tool_name]
+            if supports_rag:
+                return tool_func(query, model=model, use_rag=use_rag)
+            else:
+                return tool_func(query, model=model)
+
+    print("Tool selected by PySpur: chat_memory_tool")
+    return chat_memory_tool(query, model=model)

embedding.py

@@ -0,0 +1,47 @@
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+import os
+import zipfile
+import torch
+from transformers import AutoModel, AutoTokenizer
+import chromadb
+
+# Constants
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_NAME = "dmis-lab/biobert-base-cased-v1.1"
+DB_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "chromadb_store")
+ZIP_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "chromadb_store.zip")
+
+# Step 1: Unzip the vector store if not already present
+if not os.path.exists(os.path.join(DB_DIR, "chroma.sqlite3")):
+    print("🔓 Unzipping prebuilt ChromaDB store...")
+    with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
+        zip_ref.extractall(".")
+    print("Vector store unzipped and ready.")
+else:
+    print("Vector store already present. Skipping unzip.")
+
+# Step 2: Connect to persistent ChromaDB
+client = chromadb.PersistentClient(path=DB_DIR)
+discharge_collection = client.get_or_create_collection("discharge_notes")
+trials_collection = client.get_or_create_collection("clinical_trials")
+
+# Step 3: Load BioBERT for embedding
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)
+model.eval()
+
+# Step 4: Embedding function
+def get_embedding(text: str):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state[:, 0, :].squeeze().cpu().numpy().tolist()
+
+# Final check
+print(f"📦 ChromaDB Status:")
+print(f"  - Discharge Notes Loaded: {discharge_collection.count()}")
+print(f"  - Clinical Trials Loaded: {trials_collection.count()}")

metrics_tracker.py

@@ -0,0 +1,41 @@
+import time
+import json
+import os
+
+class MetricsTracker:
+    def __init__(self):
+        self.total_queries = 0
+        self.successful_routings = 0
+        self.failed_routings = 0
+        self.response_times = []
+
+    def record_query(self, routed_correctly: bool, response_time: float):
+        self.total_queries += 1
+        if routed_correctly:
+            self.successful_routings += 1
+        else:
+            self.failed_routings += 1
+        self.response_times.append(response_time)
+
+    def get_metrics_summary(self):
+        if self.total_queries == 0:
+            accuracy = 0.0
+            avg_response_time = 0.0
+        else:
+            accuracy = (self.successful_routings / self.total_queries) * 100
+            avg_response_time = sum(self.response_times) / self.total_queries
+
+        return {
+            "Total Queries": self.total_queries,
+            "Successful Routings": self.successful_routings,
+            "Failed Routings": self.failed_routings,
+            "Routing Accuracy (%)": round(accuracy, 2),
+            "Average Response Time (sec)": round(avg_response_time, 2)
+        }
+
+    def print_metrics_summary(self):
+        summary = self.get_metrics_summary()
+        print("\n=== Metrics Summary ===")
+        for k, v in summary.items():
+            print(f"{k}: {v}")
+        print("==========================\n")

requirements.txt

@@ -0,0 +1,28 @@
+pip>=23.2.1
+
+# NLP
+spacy==3.7.2
+en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.2/en_core_web_sm-3.7.2-py3-none-any.whl
+
+# ML / DL
+torch==2.2.0
+transformers==4.36.2
+sentence-transformers==2.2.2
+scikit-learn==1.3.2
+numpy
+pandas
+
+# PDF / File handling
+PyPDF2==3.0.1
+pdfplumber==0.10.3
+
+# Database / Vector DB
+chromadb==0.6.2
+pysqlite3-binary
+sqlalchemy>=1.4.0
+groq==0.15.0
+
+# Web / API
+streamlit==1.32.0
+google-search-results==2.4.2
+httpx==0.27.0

tool_handler.py

@@ -0,0 +1,200 @@
+import sys
+import os
+import spacy
+from groq import Groq
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'pyspur/backend/')))
+
+from embedding import discharge_collection, trials_collection, get_embedding
+
+from serpapi import GoogleSearch
+from pyspur.backend.pyspur.nodes.decorator import tool_function
+
+# Load API key from Hugging Face secret
+
+groq_api_key = os.getenv("GROQ_API_KEY")
+if not groq_api_key:
+    raise ValueError("Missing GROQ_API_KEY in environment variables.")
+
+serp_api_key = os.environ.get("SERP_API_KEY")
+if not serp_api_key:
+    raise ValueError("Missing SERP_API_KEY in environment variables.")
+
+# Initialize LLM client and spaCy
+
+client = Groq(api_key=groq_api_key)
+nlp = spacy.load("en_core_web_sm")
+
+SYMPTOM_HINTS = [
+    "chest pain", "shortness of breath", "fatigue", "dizziness",
+    "nausea", "vomiting", "palpitations", "sweating", "jaw pain",
+    "arm pain", "back pain", "tightness", "pressure in chest",
+    "arrhythmia", "tachycardia", "bradycardia", "angina",
+    "edema", "dyspnea", "syncope", "lightheadedness",
+    "ejection fraction", "myocardial infarction", "heart failure",
+    "cardiomyopathy", "cardiac arrest"
+]
+
+@tool_function(name="chat_memory_tool")
+def chat_memory_tool(memory: str, model: str = "llama-3.3-70b-versatile") -> str:
+    doc = nlp(memory)
+    found_symptoms = set(
+        keyword for chunk in doc.noun_chunks for keyword in SYMPTOM_HINTS if keyword in chunk.text.lower()
+    )
+    symptom_context = (
+        f"Previously mentioned symptoms include: {', '.join(found_symptoms)}."
+        if found_symptoms else "No clear symptoms found in memory."
+    )
+    response = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": "You are a medical assistant summarizing prior symptoms from memory."},
+            {"role": "assistant", "content": memory},
+            {"role": "user", "content": (
+                f"The patient previously reported: {memory}\n\n"
+                f"Symptoms extracted: {symptom_context}\n"
+                "Please provide a clear, concise, and helpful summary of these symptoms and suggest next steps."
+            )}
+        ]
+    )
+    return response.choices[0].message.content
+
+@tool_function(name="treatment_tool")
+def treatment_tool(query: str, model: str = "llama-3.3-70b-versatile", use_rag: bool = True) -> str:
+    try:
+        query_embedding = get_embedding(query)
+        if use_rag:
+            results = discharge_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=5,
+                include=["documents"]
+            )
+            top_docs = results['documents'][0] if results and results['documents'] else []
+            top_docs = [doc[:1500] for doc in top_docs]
+            combined_context = "\n\n".join(top_docs)
+            prompt = (
+                "You are a helpful medical assistant. Based on the following discharge notes, "
+                "recommend essential treatment.\n\n"
+                f"### Notes:\n{combined_context}\n\n### Condition:\n{query}"
+            )
+        else:
+            prompt = f"Patient condition: {query}. What treatment is recommended?"
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a medically accurate and safety-focused clinical assistant."},
+                {"role": "user", "content": prompt}
+            ]
+        )
+        return response.choices[0].message.content
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool_function(name="symptom_search_tool")
+def symptom_search_tool(symptom_description: str, model: str = "llama-3.3-70b-versatile") -> str:
+    def perform_search(query):
+        params = {
+            "engine": "google",
+            "q": f"{query} possible causes site:mayoclinic.org OR site:webmd.com OR site:nih.gov",
+            "api_key": serp_api_key
+        }
+        return GoogleSearch(params).get_dict().get("organic_results", [])
+
+    try:
+        results = perform_search(symptom_description)
+        if not results:
+            return "No reliable medical source found."
+
+        sources = []
+        snippets_with_citations = []
+        for res in results[:3]:
+            if 'snippet' in res and 'link' in res:
+                source_url = res['link']
+                domain = source_url.split("//")[-1].split("/")[0].replace("www.", "")
+                snippets_with_citations.append(f"{res['snippet']} (Source: {domain})")
+                sources.append(source_url)
+
+        search_context = "\n\n".join(snippets_with_citations)
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a medical assistant using trusted web sources to explain symptom causes."},
+                {"role": "assistant", "content": search_context},
+                {"role": "user", "content": f"What could be the cause of: {symptom_description}?"}
+            ]
+        )
+
+        bulleted_sources = "\n".join(f"- {url}" for url in sources)
+        return response.choices[0].message.content + "\n\n**Sources:**\n" + bulleted_sources
+
+    except Exception as e:
+        return f"Search error: {str(e)}"
+
+@tool_function(name="trial_matcher_tool")
+def trial_matcher_tool(discharge_note: str, model: str = "llama-3.3-70b-versatile", use_rag: bool = True) -> str:
+    try:
+        query_embedding = get_embedding(discharge_note)
+        results = trials_collection.query(
+            query_embeddings=[query_embedding],
+            n_results=3,
+            include=["documents", "metadatas"]
+        )
+        if not results.get('documents') or not results['documents'][0]:
+            return "No matching clinical trials were found for the provided note."
+
+        summaries = []
+        for i, (doc, meta) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
+            nct_id = meta.get("NCT ID") or "Unknown ID"
+            truncated_doc = doc.strip()[:1500]
+            if use_rag:
+                summary_prompt = (
+                    f"You are a clinical assistant reviewing a matched clinical trial.\n"
+                    f"Summarize the trial using **bullet points only** for the following fields:\n"
+                    f"- NCT ID\n- Study Title\n- Conditions\n- Inclusion Criteria\n- Exclusion Criteria\n\n"
+                    f"Use bullets under each field. Maintain a clean format. Respond only with the summary.\n\n"
+                    f"Trial Description:\nNCT ID: {nct_id}\n{truncated_doc}"
+                )
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=[
+                        {"role": "system", "content": "You are a medically precise clinical research assistant."},
+                        {"role": "user", "content": summary_prompt}
+                    ]
+                )
+                summaries.append(f"### Trial {i+1}:\n{response.choices[0].message.content}")
+            else:
+                summaries.append(f"### Trial {i+1}:\nNCT ID: {nct_id}\n\n{truncated_doc}")
+
+        return "\n\n---\n\n".join(summaries)
+
+    except Exception as e:
+        return f"Error during trial matching: {str(e)}"
+
+# Tool routing via keyword logic
+
+TOOL_ROUTER = {
+    "symptom": ("symptom_search_tool", False),
+    "treatment": ("treatment_tool", True),
+    "trial": ("trial_matcher_tool", True)
+}
+
+TOOL_FUNCTIONS = {
+    "chat_memory_tool": chat_memory_tool,
+    "treatment_tool": treatment_tool,
+    "symptom_search_tool": symptom_search_tool,
+    "trial_matcher_tool": trial_matcher_tool
+}
+
+def run_tool(query: str, model: str, use_rag: bool) -> str:
+    for keyword, (tool_name, supports_rag) in TOOL_ROUTER.items():
+        if keyword in query.lower():
+            print(f"Tool selected by PySpur: {tool_name}")
+            tool_func = TOOL_FUNCTIONS[tool_name]
+            if supports_rag:
+                return tool_func(query, model=model, use_rag=use_rag)
+            else:
+                return tool_func(query, model=model)
+
+    print("Tool selected by PySpur: chat_memory_tool")
+    return chat_memory_tool(query, model=model)