Add search_kaggle_datasets, download_kaggle_dataset and authorize the use of datasci libs

Gradio_UI.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+# pandas transformers sklearn smolagents matplotlib seaborn
 import mimetypes
 import os
 import re

app.py

@@ -4,8 +4,10 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
+# from typing import Optional
 
 from kaggle.api.kaggle_api_extended import KaggleApi
+import os
 
 from Gradio_UI import GradioUI
 
@@ -21,7 +23,7 @@ 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:
+def search_kaggle_datasets(search_term:str, kaggle_username:str = None, kaggle_key:str = None, max_results:int = 10)-> list[dict[str]]:
     """Search for datasets on Kaggle based on a search term.
     Args:
         search_term: The term to search for.
@@ -45,8 +47,11 @@ def search_kaggle_datasets(search_term:str, kaggle_username=None:str, kaggle_key
         # Use the default kaggle.json file if no credentials are provided
         return 'Error in searching Kaggle datasets: No username or key provided.'
 
-    api.authenticate()
-
+    try:
+        api.authenticate()
+    except Exception as e:
+        return f"Error authenticating with Kaggle: {str(e)}"
+    
     # Search for datasets
     datasets = api.dataset_list(search=search_term)
 
@@ -71,22 +76,54 @@ def search_kaggle_datasets(search_term:str, kaggle_username=None:str, kaggle_key
 
     return results
 
-
 @tool
-def get_current_time_in_timezone(timezone: str) -> str:
-    """A tool that fetches the current local time in a specified timezone.
+def download_kaggle_dataset(
+    dataset_ref: str,
+    download_path: str,
+    kaggle_username: str = None,
+    kaggle_key: str = None,
+    unzip: bool = True
+) -> str:
+    """Download a dataset from Kaggle.
     Args:
-        timezone: A string representing a valid timezone (e.g., 'America/New_York').
+        dataset_ref: The reference of the dataset (e.g., "username/dataset-name").
+        download_path: The directory where the dataset will be downloaded.
+        kaggle_username: Your Kaggle username (from kaggle.json).
+        kaggle_key: Your Kaggle API key (from kaggle.json).
+        unzip: Whether to unzip the dataset after downloading. Default is True.
     """
+    # 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
+        pass
+
     try:
-        # Create timezone object
-        tz = pytz.timezone(timezone)
-        # Get current time in that timezone
-        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
-        return f"The current local time in {timezone} is: {local_time}"
+        api.authenticate()
     except Exception as e:
-        return f"Error fetching time for timezone '{timezone}': {str(e)}"
+        return f"Error authenticating with Kaggle: {str(e)}"
+
+    # Ensure the download path exists
+    os.makedirs(download_path, exist_ok=True)
+
+    # Download the dataset
+    api.dataset_download_files(dataset_ref, path=download_path, unzip=unzip)
+
+    # Clean up the temporary kaggle.json file if it was created
+    if kaggle_username and kaggle_key:
+        os.remove(kaggle_json_path)
 
+    return f"Dataset '{dataset_ref}' downloaded to '{download_path}'."
 
 final_answer = FinalAnswerTool()
 
@@ -109,14 +146,15 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[final_answer, search_kaggle_datasets, download_kaggle_dataset], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,
     planning_interval=None,
     name=None,
     description=None,
-    prompt_templates=prompt_templates
+    prompt_templates=prompt_templates,
+    additional_authorized_imports=['pandas', 'matplotlib', 'seaborn'],
 )