Update app.py

app.py

@@ -7,36 +7,44 @@ from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
 # Custom Tool to fetch datasets related to body parts or imaging types
+
 @tool
-def my_custom_tool(arg1: str, arg2: int) -> str:
+def medsearch_tool(arg1: str, arg2: int) -> str:
     """
-    Fetch datasets related to a medical keyword such as body parts (e.g., skin, brain) or imaging modalities (e.g., MRI, CT).
+    Search and retrieve publicly available medical datasets from Hugging Face based on any medical-related keyword.
 
     Args:
-        arg1: A medical keyword such as 'skin', 'brain', 'XRay', or 'ultrasound'.
-        arg2: The maximum number of datasets to return.
+        arg1: A keyword related to medical data (e.g., 'cancer', 'diabetes', 'CT scan', 'radiology', 'dermoscopy').
+        arg2: The maximum number of datasets to retrieve.
 
     Returns:
-        A formatted list of dataset names matching the keyword from Hugging Face.
+        A formatted list of dataset names and summaries matching the search query,
+        or a message stating that no datasets were found.
     """
     try:
-        query = arg1.lower()
+        keyword = arg1.strip()
         limit = int(arg2)
 
         response = requests.get(
-            f"https://huggingface.co/api/datasets?search={query}&limit={limit}"
+            f"https://huggingface.co/api/datasets?search={keyword}&limit={limit}"
         )
         response.raise_for_status()
         datasets = response.json()
 
-        if not datasets:
-            return f"No datasets found for '{arg1}'."
+        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}")
 
-        results = [f"- {ds.get('id', 'Unknown')}" for ds in datasets[:limit]]
-        return f"Datasets related to '{arg1}':\n" + "\n".join(results)
+        return f"Medical datasets related to '{keyword}':\n" + "\n\n".join(results)
 
     except Exception as e:
-        return f"Error fetching datasets for '{arg1}': {str(e)}"
+        return f"Error searching medical datasets for '{arg1}': {str(e)}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -76,7 +84,7 @@ with open("prompts.yaml", 'r') as stream:
 # Create the agent
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_current_time_in_timezone, my_custom_tool],  # add your tools here
+    tools=[final_answer, get_current_time_in_timezone, medsearch_tool],  # add your tools here
     max_steps=6,
     verbosity_level=1,
     grammar=None,