|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
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.") |
|
|
|
|
|
|
|
|
|
|
|
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_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): |
|
|
for keyword, (tool_name, supports_rag) in TOOL_ROUTER.items(): |
|
|
if keyword in query.lower(): |
|
|
print(f"[ROUTER] Tool selected: {tool_name}") |
|
|
print(f"[ROUTER] Model: {model} | RAG: {use_rag}") |
|
|
|
|
|
tool_func = TOOL_FUNCTIONS[tool_name] |
|
|
|
|
|
if supports_rag: |
|
|
response = tool_func(query, model=model, use_rag=use_rag) |
|
|
else: |
|
|
response = tool_func(query, model=model) |
|
|
|
|
|
return response, tool_name |
|
|
|
|
|
|
|
|
tool_name = "chat_memory_tool" |
|
|
print(f"[ROUTER] Tool selected: {tool_name}") |
|
|
print(f"[ROUTER] Model: {model} | RAG: {use_rag}") |
|
|
|
|
|
response = chat_memory_tool(query, model=model) |
|
|
return response, tool_name |
