Updated app.py tool

app.py

@@ -9,7 +9,7 @@ from Gradio_UI import GradioUI
 # Custom Tool to fetch datasets related to body parts or imaging types
 
 @tool
-def medsearch_tool(arg1: str, arg2: int) -> str:
+def my_custom_tool(arg1: str, arg2: int) -> str:
     """
     Search and retrieve publicly available medical datasets from Hugging Face based on any medical-related keyword.
 
@@ -18,30 +18,38 @@ def medsearch_tool(arg1: str, arg2: int) -> str:
         arg2: The maximum number of datasets to retrieve.
 
     Returns:
-        A formatted list of dataset names and summaries matching the search query,
-        or a message stating that no datasets were found.
+        A list of dataset names matching the search query, or a message stating that no datasets were found
+        or that the keyword is not medically relevant.
     """
     try:
-        keyword = arg1.strip()
+        keyword = arg1.strip().lower()
         limit = int(arg2)
 
+        # ✅ Define a basic list of medically relevant terms
+        medical_terms = [
+            "skin", "brain", "cancer", "tumor", "mri", "ct", "xray", "ultrasound",
+            "diabetes", "pneumonia", "covid", "lesion", "radiology", "pathology",
+            "lung", "chest", "abdomen", "spine", "bone", "stroke", "eczema", "melanoma",
+            "eye", "retina", "dermoscopy", "cardiology", "infection", "biopsy", "tooth",
+            "toothache", "dental", "ear", "wrist", "hand", "leg", "arm", "heart"
+        ]
+
+        # ✅ Check if the keyword looks medically relevant
+        if not any(term in keyword for term in medical_terms):
+            return f"'{arg1}' does not appear to be a medical term."
+
+        # ✅ Proceed to fetch datasets
         response = requests.get(
             f"https://huggingface.co/api/datasets?search={keyword}&limit={limit}"
         )
         response.raise_for_status()
         datasets = response.json()
 
-        if not datasets or len(datasets) == 0:
-            return f"No medical datasets found for '{keyword}'."
-
-        results = []
-        for ds in datasets[:limit]:
-            dataset_id = ds.get('id', 'Unknown')
-            card_data = ds.get('cardData', {})
-            description = card_data.get('summary', '').strip() or "No description available."
-            results.append(f"- {dataset_id}: {description}")
+        if not datasets:
+            return f"No medical datasets found for '{arg1}'."
 
-        return f"Medical datasets related to '{keyword}':\n" + "\n\n".join(results)
+        results = [f"- {ds.get('id', 'Unknown')}" for ds in datasets[:limit]]
+        return f"Medical datasets related to '{arg1}':\n" + "\n".join(results)
 
     except Exception as e:
         return f"Error searching medical datasets for '{arg1}': {str(e)}"
@@ -84,7 +92,7 @@ with open("prompts.yaml", 'r') as stream:
 # Create the agent
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_current_time_in_timezone, medsearch_tool],  # add your tools here
+    tools=[final_answer, get_current_time_in_timezone, my_custom_tool],  # add your tools here
     max_steps=6,
     verbosity_level=1,
     grammar=None,