Update tool_handler.py

tool_handler.py

@@ -3,18 +3,15 @@ import os
 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 groq import Client
 from serpapi import GoogleSearch
 import spacy
-from pyspur import SpurAgent
 from pyspur.nodes.decorator import tool_function
 
 # Initialize LLM client and spaCy
 client = Client(api_key="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")
 nlp = spacy.load("en_core_web_sm")
 
-# Symptom keywords
 SYMPTOM_HINTS = [
     "chest pain", "shortness of breath", "fatigue", "dizziness",
     "nausea", "vomiting", "palpitations", "sweating", "jaw pain",
@@ -25,9 +22,8 @@ SYMPTOM_HINTS = [
     "cardiomyopathy", "cardiac arrest"
 ]
 
-# Tool 1: Chat Memory Symptom Reasoner
-@tool_function
-def chat_memory_tool(memory: str, model="llama-3.3-70b-versatile") -> str:
+@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()
@@ -50,12 +46,10 @@ def chat_memory_tool(memory: str, model="llama-3.3-70b-versatile") -> str:
     )
     return response.choices[0].message.content
 
-# Tool 2: Treatment Recommender
-@tool_function
-def treatment_tool(query: str, model="llama-3.3-70b-versatile", use_rag=True) -> str:
+@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],
@@ -65,7 +59,6 @@ def treatment_tool(query: str, model="llama-3.3-70b-versatile", use_rag=True) ->
             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"
@@ -86,9 +79,8 @@ def treatment_tool(query: str, model="llama-3.3-70b-versatile", use_rag=True) ->
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Tool 3: Symptom Cause Analyzer
-@tool_function
-def symptom_search_tool(symptom_description: str, model="llama-3.3-70b-versatile") -> str:
+@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",
@@ -104,7 +96,6 @@ def symptom_search_tool(symptom_description: str, model="llama-3.3-70b-versatile
 
         sources = []
         snippets_with_citations = []
-
         for res in results[:3]:
             if 'snippet' in res and 'link' in res:
                 source_url = res['link']
@@ -113,7 +104,6 @@ def symptom_search_tool(symptom_description: str, model="llama-3.3-70b-versatile
                 sources.append(source_url)
 
         search_context = "\n\n".join(snippets_with_citations)
-
         response = client.chat.completions.create(
             model=model,
             messages=[
@@ -129,18 +119,15 @@ def symptom_search_tool(symptom_description: str, model="llama-3.3-70b-versatile
     except Exception as e:
         return f"Search error: {str(e)}"
 
-# Tool 4: Clinical Trial Matcher
-@tool_function
-def trial_matcher_tool(discharge_note: str, model="llama-3.3-70b-versatile", use_rag=True) -> str:
+@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."
 
@@ -148,7 +135,6 @@ def trial_matcher_tool(discharge_note: str, model="llama-3.3-70b-versatile", use
         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"
@@ -157,7 +143,6 @@ def trial_matcher_tool(discharge_note: str, model="llama-3.3-70b-versatile", use
                     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=[
@@ -174,10 +159,15 @@ def trial_matcher_tool(discharge_note: str, model="llama-3.3-70b-versatile", use
     except Exception as e:
         return f"Error during trial matching: {str(e)}"
 
-# Register tools in PySpur Agent
-agent = SpurAgent(tools=[
-    chat_memory_tool,
-    treatment_tool,
-    symptom_search_tool,
-    trial_matcher_tool
-])
+# Tool routing via keyword logic
+TOOL_ROUTER = {
+    "symptom": symptom_search_tool,
+    "treatment": treatment_tool,
+    "trial": trial_matcher_tool
+}
+
+def run_tool(query: str, model: str, use_rag: bool) -> str:
+    for keyword, tool_func in TOOL_ROUTER.items():
+        if keyword in query.lower():
+            return tool_func(query, model=model, use_rag=use_rag)
+    return chat_memory_tool(query, model=model)