Add search_kaggle_datasets

app.py

@@ -1,10 +1,12 @@
-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 kaggle.api.kaggle_api_extended import KaggleApi
+
 from Gradio_UI import GradioUI
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
@@ -18,6 +20,58 @@ def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return
     """
     return "What magic will you build ?"
 
+@tool
+def search_kaggle_datasets(search_term:str, kaggle_username=None:str, kaggle_key=None:str, max_results:int)-> str:
+    """Search for datasets on Kaggle based on a search term.
+    Args:
+        search_term: The term to search for.
+        kaggle_username: Your Kaggle username.
+        kaggle_key: Your Kaggle API key.
+        max_results: Maximum number of results to return.
+    """
+     # Initialize the Kaggle API
+    api = KaggleApi()
+
+    # Authenticate using provided credentials
+    if kaggle_username and kaggle_key:
+        # Create a temporary kaggle.json file
+        kaggle_json_content = f'{{"username":"{kaggle_username}","key":"{kaggle_key}"}}'
+        kaggle_json_path = os.path.expanduser("~/.kaggle/kaggle.json")
+        os.makedirs(os.path.dirname(kaggle_json_path), exist_ok=True)
+        with open(kaggle_json_path, "w") as f:
+            f.write(kaggle_json_content)
+        os.chmod(kaggle_json_path, 0o600)  # Set permissions to read/write for the owner only
+    else:
+        # Use the default kaggle.json file if no credentials are provided
+        return 'Error in searching Kaggle datasets: No username or key provided.'
+
+    api.authenticate()
+
+    # Search for datasets
+    datasets = api.dataset_list(search=search_term)
+
+    # Limit the number of results
+    datasets = datasets[:max_results]
+
+    # Extract relevant information
+    results = []
+    for dataset in datasets:
+        dataset_info = api.dataset_view(dataset)
+        results.append({
+            'title': dataset_info['title'],
+            'url': f"https://www.kaggle.com/{dataset_info['ref']}",
+            'size': dataset_info['size'],
+            'files': dataset_info['files'],
+            'last_updated': dataset_info['lastUpdated']
+        })
+
+    # Clean up the temporary kaggle.json file if it was created
+    if kaggle_username and kaggle_key:
+        os.remove(kaggle_json_path)
+
+    return results
+
+
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """A tool that fetches the current local time in a specified timezone.