Update app.py

app.py

@@ -1,167 +1,127 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
 from Gradio_UI import GradioUI
 
-# Below is an example of a tool that does nothing. Amaze us with your creativity !
-# @tool
-# def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
-#     #Keep this format for the description / args / args description but feel free to modify the tool
-#     """A tool that does nothing yet 
-#     Args:
-#         arg1: the first argument
-#         arg2: the second argument
-#     """
-#     return "What magic will you build ?"
-
-
-def __init__(self):
-    self.patient_data = {}
-
-@tool
-def collect_symptoms(patient_id: str, symptoms: str) -> str:
-    """
-    Nurse tool to collect symptoms and ask follow-up questions.
-
-    Args:
-        patient_id: Unique identifier for the patient.
-        symptoms: Initial symptoms provided by the patient.
-
-    Returns:
-        Follow-up questions based on the symptoms.
-    """
-    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 taking any medications?",
-            "Have you experienced these symptoms before?",
-            "Have you had any recent illnesses?",
-            "Have you noticed any other unusual changes?",
-            "What is your medical history related to these symptoms?"
-            
-        ]
-
-        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)}"
-
-@tool
-def diagnose_patient(patient_id: str) -> str:
-    """
-    Doctor tool to diagnose the patient based on symptoms.
-    
-    Args:
-        patient_id: Unique identifier for the patient.
-        
-    Returns:
-        Diagnosis and recommendations.
-    """
-    try:
-        if patient_id not in CodeAgent.patient_data:
-            raise ValueError("No symptoms found. Nurse must collect symptoms first.")
-
-        symptoms = CodeAgent.patient_data[patient_id]["symptoms"]
-        diagnosis = CodeAgent.fetch_diagnosis(symptoms)
-        medication = CodeAgent.fetch_medication(symptoms)
-        advice = CodeAgent.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(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(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(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()
+class MedicalAgent:
+    def __init__(self):
+        self.patient_data = {}
+        self.search_tool = DuckDuckGoSearchTool()
+
+    @tool
+    def collect_symptoms(self, patient_id: str, symptoms: str) -> str:
+        """
+        Nurse tool to collect symptoms and ask follow-up questions.
+        Args:
+            patient_id: Unique identifier for the patient.
+            symptoms: Initial symptoms provided by the patient.
+        Returns:
+            Follow-up questions based on the symptoms.
+        """
+        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 taking any medications?",
+                "Have you experienced these symptoms before?",
+                "Have you had any recent illnesses?",
+                "Have you noticed any other unusual changes?",
+                "What is your medical history related to these symptoms?"
+            ]
+
+            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)}"
+
+    @tool
+    def diagnose_patient(self, patient_id: str) -> str:
+        """
+        Doctor tool to diagnose the patient based on symptoms.
+        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.
+        """
+        try:
+            if not symptoms:
+                raise ValueError("Symptoms must be provided.")
+            search_query = f"A detailed medical diagnosis for these symptoms: {symptoms}"
+            return self.search_tool(search_query)
+        except Exception as e:
+            return f"Error fetching diagnosis: {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.
+        """
+        try:
+            if not symptoms:
+                raise ValueError("Symptoms must be provided.")
+            search_query = f"Recommended medications for these symptoms: {symptoms}"
+            return self.search_tool(search_query)
+        except Exception as e:
+            return f"Error fetching medication: {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.
+        """
+        try:
+            if not symptoms:
+                raise ValueError("Symptoms must be provided.")
+            search_query = f"Treatment advice and recommendations for these symptoms: {symptoms}"
+            return self.search_tool(search_query)
+        except Exception as e:
+            return f"Error fetching treatment advice: {str(e)}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -170,45 +130,53 @@ def get_current_time_in_timezone(timezone: str) -> str:
         timezone: A string representing a valid timezone (e.g., 'America/New_York').
     """
     try:
-        # Create timezone object
         tz = pytz.timezone(timezone)
-        # Get current time in that timezone
         local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
         return f"The current local time in {timezone} is: {local_time}"
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-
-final_answer = FinalAnswerTool()
-
-# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
-# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 
-
-model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
-)
-
-
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
-    
-agent = CodeAgent(
-    model=model,
-    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="MedicalDoctor",
-    description="A tool that simulates the roles of a medical doctor and a nurse in gathering symptoms, diagnosing, and providing recommendations",
-    prompt_templates=prompt_templates
-)
-
-
-GradioUI(agent).launch()
+def initialize_agent():
+    try:
+        with open("prompts.yaml", 'r') as stream:
+            prompt_templates = yaml.safe_load(stream)
+    except FileNotFoundError:
+        print("Warning: prompts.yaml not found. Using default templates.")
+        prompt_templates = {}
+
+    final_answer = FinalAnswerTool()
+    medical_agent = MedicalAgent()
+
+    # If the primary model is overloaded, uncomment the following alternative endpoint
+    model_id = 'https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
+    model = HfApiModel(
+        max_tokens=2096,
+        temperature=0.5,
+        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
+        custom_role_conversions=None,
+    )
+
+    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+    agent = CodeAgent(
+        model=model,
+        tools=[
+            final_answer,
+            medical_agent.collect_symptoms,
+            medical_agent.diagnose_patient,
+            get_current_time_in_timezone
+        ],
+        max_steps=6,
+        verbosity_level=1,
+        grammar=None,
+        planning_interval=None,
+        name="MedicalDoctor",
+        description="A tool that simulates the roles of a medical doctor and a nurse in gathering symptoms, diagnosing, and providing recommendations",
+        prompt_templates=prompt_templates
+    )
+
+    return agent
+
+if __name__ == "__main__":
+    agent = initialize_agent()
+    GradioUI(agent).launch()