Update app.py

app.py

@@ -255,21 +255,76 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
         arg2: The maximum number of datasets to retrieve.
 
     Returns:
-        A list of dataset names matching the search query, or a message stating that no datasets were found.
+        A numbered list (top N) of dataset names matching the search query.
     """
     try:
         keyword = arg1.strip().lower()
         limit = int(arg2)
 
+       # Define a comprehensive list of medically relevant terms
         medical_terms = [
-            "skin", "brain", "lung", "chest", "abdomen", "spine", "bone", "heart", "liver",
-            "radiology", "pathology", "oncology", "dermatology", "mri", "ct", "xray", "ultrasound",
-            "cancer", "tumor", "melanoma", "eczema", "thyroid"
+            # 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
+            "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"
         ]
 
+
         if not any(term in keyword for term in medical_terms):
             return f"No medical datasets found for '{arg1}'. Please try another medical term."
 
+        # Query Hugging Face API
         try:
             response = requests.get(
                 f"https://huggingface.co/api/datasets?search={keyword}&limit={limit}",
@@ -278,13 +333,17 @@ def my_custom_tool(arg1: str, arg2: int) -> str:
             response.raise_for_status()
             datasets = response.json()
         except Exception:
+            # Offline fallback
             datasets = [{"id": f"example/{keyword}-dataset-{i+1}"} for i in range(limit)]
 
         if not datasets:
             return f"No datasets found for '{arg1}'."
 
-        results = [f"- {ds.get('id', 'Unknown')}" for ds in datasets[:limit]]
-        return f"Medical datasets related to '{arg1}':\n" + "\n".join(results)
+        # Format results neatly with numbered bullets
+        formatted = "\n".join(
+            [f"- Dataset {i+1}: {ds.get('id', 'Unknown')}" for i, ds in enumerate(datasets[:limit])]
+        )
+        return f"Medical datasets related to '{arg1}':\n{formatted}"
 
     except Exception as e:
         return f"Error searching medical datasets for '{arg1}': {str(e)}"