Update app.py

updated version with more tools

app.py

@@ -18,39 +18,148 @@ from Gradio_UI import GradioUI
 #     """
 #     return "What magic will you build ?"
 
+
+def __init__(self):
+    self.patient_data = {}
+
 @tool
-def med_tool(role: str, symptoms: str) -> str:
-    # It's important to specify the return type
-    """A tool for symptom collection (nurse) and diagnosis (doctor).
-    
+def collect_symptoms(self, patient_id: str, symptoms: str) -> str:
+    """
+    Nurse tool to collect symptoms and ask follow-up questions.
+
     Args:
-        role: The role of the healthcare provider ('nurse' or 'doctor')
-        symptoms: The symptoms described by the patient
-    
+        patient_id: Unique identifier for the patient.
+        symptoms: Initial symptoms provided by the patient.
+
     Returns:
-        A response based on the role provided.
+        Follow-up questions based on the symptoms.
     """
-    
-    if role.lower() == "nurse":
-        follow_up_questions = [
+    try:
+        if not patient_id or not symptoms:
+            raise ValueError("Patient ID and symptoms must be provided.")
+
+        self.patient_data[patient_id] = {"symptoms": symptoms, "additional_info": {}}
+
+        questions = [
             "How long have you had these symptoms?",
             "Do you have any allergies?",
-            "Are you on any medication?",
+            "Are you taking any medications?",
+            "Have you experienced these symptoms before?",
             "Have you had any recent illnesses?",
-            "Have you noticed any other unusual changes?"
+            "Have you noticed any other unusual changes?",
+            "What is your medical history related to these symptoms?"
+            
         ]
-        return (f"Nurse: I have recorded the symptoms: {symptoms}.\n"
-                f"Here are some follow-up questions:\n- " + "\n- ".join(follow_up_questions))
 
-    elif role.lower() == "doctor":
-        # Delegating diagnosis to an external tool (e.g., AI search, RAG)
-        return (f"Doctor: Based on the symptoms: {symptoms}, I will now analyze the data.\n"
-                f"Fetching diagnosis, medication, and treatment recommendations from the medical database...")
+        return f"Nurse: I have noted the symptoms ({symptoms}). Here are follow-up questions:\n" + "\n".join(questions)
+
+    except ValueError as e:
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
 
-    else:
-        return "Invalid role! Please specify either 'nurse' or 'doctor'."
+@tool
+def diagnose_patient(self, patient_id: str) -> str:
+    """
+    Doctor tool to diagnose the patient based on symptoms.
 
-final_answer = FinalAnswerTool()
+    Args:
+        patient_id: Unique identifier for the patient.
+
+    Returns:
+        Diagnosis and recommendations.
+    """
+    try:
+        if patient_id not in self.patient_data:
+            raise ValueError("No symptoms found. Nurse must collect symptoms first.")
+
+        symptoms = self.patient_data[patient_id]["symptoms"]
+        diagnosis = self.fetch_diagnosis(symptoms)
+        medication = self.fetch_medication(symptoms)
+        advice = self.fetch_treatment_advice(symptoms)
+
+        return (f"Doctor: Based on the symptoms: {symptoms},\n"
+                f"Diagnosis: {diagnosis}\n"
+                f"Medication: {medication}\n"
+                f"Advice: {advice}")
+
+    except ValueError as e:
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
+
+@tool
+def fetch_diagnosis(self, symptoms: str) -> str:
+    """
+    AI tool to retrieve a diagnosis based on symptoms.
+
+    Args:
+        symptoms: The symptoms provided by the patient.
+
+    Returns:
+        Detailed Diagnosis information.
+    """
+    search_query = f" A detailed medical diagnosis for this {symptoms}"
+    return search_tool(search_query)
+    try:
+        if not symptoms:
+            raise ValueError("Symptoms must be provided.")
+
+        return f"AI Diagnosis: Potential cause of {symptoms} found."
+
+    except ValueError as e:
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
+
+@tool
+def fetch_medication(self, symptoms: str) -> str:
+    """
+    AI tool to suggest medications based on symptoms.
+
+    Args:
+        symptoms: The symptoms provided by the patient.
+
+    Returns:
+        Suggested medication.
+    """
+    search_query = f" A detailed suggested medication for this {symptoms}"
+    return search_tool(search_query)
+    try:
+        if not symptoms:
+            raise ValueError("Symptoms must be provided.")
+
+        return f"Suggested medication for {symptoms} found."
+
+    except ValueError as e:
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
+
+@tool
+def fetch_treatment_advice(self, symptoms: str) -> str:
+    """
+    AI tool to provide treatment recommendations.
+
+    Args:
+        symptoms: The symptoms provided by the patient.
+
+    Returns:
+        Recommended treatment and advice.
+    """
+    search_query = f" A detailed recommended treatment and advice for this {symptoms}"
+    return search_tool(search_query)
+    try:
+        if not symptoms:
+            raise ValueError("Symptoms must be provided.")
+
+        return f"Treatment advice for {symptoms} retrieved."
+
+    except ValueError as e:
+        return f"Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
+        
 search_tool = DuckDuckGoSearchTool()
 
 @tool
@@ -90,13 +199,13 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, med_tool], ## add your tools here (don't remove final answer)
+    tools=[final_answer, collect_symptoms, diagnose_patient, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,
     planning_interval=None,
-    name=None,
-    description=None,
+    name="MedicalDoctor",
+    description"=The role of the healthcare provider ('nurse' or 'doctor')",
     prompt_templates=prompt_templates
 )