Update app.py

app.py

@@ -116,71 +116,26 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
 
     Returns:
         A list of dataset names matching the search query, or a message stating that no datasets were found.
-        If network access is restricted, simulated dataset results are returned instead.
     """
     try:
         keyword = arg1.strip().lower()
         limit = int(arg2)
 
-        # Define a basic list of medically relevant terms
+        # Define a list of medical terms
         medical_terms = [
-            # Anatomy / Body Parts
             "skin", "brain", "lung", "chest", "abdomen", "spine", "bone", "heart", "liver", "kidney",
-            "bladder", "stomach", "colon", "rectum", "esophagus", "pancreas", "breast", "ear", "eye", 
-            "retina", "tooth", "teeth", "tongue", "jaw", "neck", "wrist", "hand", "leg", "arm", "shoulder", "pelvis",
-        
-            # Diseases / Conditions
-            "cancer", "tumor", "stroke", "diabetes", "pneumonia", "covid", "asthma", "eczema", "melanoma",
-            "hypertension", "alzheimer", "parkinson", "arthritis", "scoliosis", "epilepsy", "glaucoma",
-            "ulcer", "hepatitis", "leukemia", "lymphoma", "tuberculosis", "anemia", "obesity", "depression",
-            "anxiety", "bipolar", "autism", "adhd", "ptsd", "psychosis", "schizophrenia",
-        
-            # Imaging Modalities
-            "mri", "ct", "xray", "x-ray", "ultrasound", "pet", "fmri", "mammo", "angiography", "radiography",
-            "echocardiogram", "spect", "dermoscopy", "colonoscopy", "endoscopy", "biopsy", "histopathology",
-        
-            # Medical Specialties
+            "bladder", "stomach", "colon", "rectum", "esophagus", "pancreas", "breast", "ear", "eye",
             "radiology", "pathology", "oncology", "cardiology", "neurology", "dermatology", "dentistry",
             "ophthalmology", "urology", "orthopedics", "gastroenterology", "pulmonology", "nephrology",
             "psychiatry", "pediatrics", "geriatrics", "infectious disease",
-        
-            # Symptoms / Signs
-            "lesion", "infection", "fever", "pain", "inflammation", "rash", "headache", "swelling", 
-            "cough", "seizure", "dizziness", "vomiting", "diarrhea", "nausea", "fatigue", "itching",
-        
-            # Common Specific Diseases
-            "breast cancer", "prostate cancer", "lung cancer", "skin cancer", "colon cancer", 
-            "brain tumor", "liver cancer", "cervical cancer", "bladder cancer", "thyroid cancer",
-        
-            # Procedures / Interventions
-            "surgery", "chemotherapy", "radiation", "transplant", "dialysis", "intubation", "stenting",
-            "ventilation", "vaccination", "anesthesia", "rehabilitation", "prosthetics", "orthotics",
-        
-            # Lab Tests / Biomarkers
-            "blood test", "cbc", "glucose", "hemoglobin", "cholesterol", "biomarker", "urinalysis",
-            "pcr", "serology", "antibody", "antigen",
-        
-            # Clinical Settings / Roles
-            "icu", "hospital", "emergency", "clinical notes", "nursing", "physician", "patient",
-            "medical record", "electronic health record", "ehr", "vitals",
-        
-            # Age-based Terms
-            "pediatric", "neonatal", "infant", "child", "adolescent", "geriatrics", "elderly",
-        
-            # Epidemiology / Public Health
-            "epidemiology", "prevalence", "incidence", "mortality", "public health", "health disparity",
-            "risk factor", "social determinant",
-        
-            # Pharmacology / Medications
-            "drug", "medication", "pharmacology", "side effect", "adverse event", "dose", "tablet",
-            "vaccine", "clinical trial", "placebo"
+            "mri", "ct", "xray", "x-ray", "ultrasound", "pet", "fmri", "mammo", "angiography", "radiography",
+            "cancer", "tumor", "stroke", "diabetes", "melanoma", "eczema", "asthma", "thyroid"
         ]
 
-        # Check if keyword is in known medical terms
         if not any(term in keyword for term in medical_terms):
-            return f"No medical datasets found for '{arg1}'."
+            return f"No medical datasets found for '{arg1}'. Please try another medical term."
 
-        # Try fetching datasets from Hugging Face, fallback to mock data if blocked
+        # Try online query to Hugging Face
         try:
             response = requests.get(
                 f"https://huggingface.co/api/datasets?search={keyword}&limit={limit}",
@@ -188,15 +143,14 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
             )
             response.raise_for_status()
             datasets = response.json()
-        except requests.exceptions.RequestException:
-            # Simulate results if network access is blocked
+        except Exception:
+            # Network-restricted fallback
             datasets = [{"id": f"example/{keyword}-dataset-{i+1}"} for i in range(limit)]
 
-        # Return message if no datasets found
+        # Return formatted list
         if not datasets:
-            return f"No medical datasets found for '{arg1}'."
+            return f"No datasets found for '{arg1}'."
 
-        # Collect and return dataset names
         results = [f"- {ds.get('id', 'Unknown')}" for ds in datasets[:limit]]
         return f"Medical datasets related to '{arg1}':\n" + "\n".join(results)
 
@@ -206,6 +160,7 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
 
 
 
+
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """
@@ -250,15 +205,30 @@ model = InferenceClientModel(
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-# Create the agent
+# # Create the agent
+# agent = CodeAgent(
+#     model=model,
+#     tools=[final_answer, get_current_time_in_timezone, my_custom_tool],
+#     max_steps=6,
+#     verbosity_level=2,
+#     planning_interval=None,
+#     name=None,
+#     description=None,
+#     prompt_templates=prompt_templates
+# )
+
 agent = CodeAgent(
     model=model,
     tools=[final_answer, get_current_time_in_timezone, my_custom_tool],
     max_steps=6,
-    verbosity_level=2,
+    verbosity_level=1,
     planning_interval=None,
-    name=None,
-    description=None,
+    name="MedDataSearchAgent",
+    description=(
+        "An intelligent agent that searches Hugging Face datasets related to "
+        "medical conditions, body parts, and imaging modalities. "
+        "Use 'my_custom_tool' whenever the user requests medical data or datasets."
+    ),
     prompt_templates=prompt_templates
 )