Update example.ipynb

Author: ludomitch

tutorial/example.ipynb

@@ -36,7 +36,7 @@
    "outputs": [],
    "source": [
     "def setup_data_analysis_env(\n",
-    "    query: str, dataset_folder: Path, language: NBLanguage = NBLanguage.PYTHON\n",
+    "    query: str, dataset: Path, language: NBLanguage = NBLanguage.PYTHON\n",
     "):\n",
     "    # Hash the task to get a unique identifier\n",
     "    task_hash = hashlib.sha256(query.encode()).hexdigest()\n",
@@ -46,11 +46,14 @@
     "    trajectory_path.mkdir(parents=True, exist_ok=True)\n",
     "    nb_path = trajectory_path / NBEnvironment.NOTEBOOK_NAME\n",
     "    # Copy task data to trajectory path\n",
-    "    for item in dataset_folder.iterdir():\n",
-    "        if item.is_file():\n",
-    "            shutil.copy2(item, trajectory_path)\n",
-    "        elif item.is_dir():\n",
-    "            shutil.copytree(item, trajectory_path / item.name, dirs_exist_ok=True)\n",
+    "    if dataset.is_dir():\n",
+    "        for item in dataset.iterdir():\n",
+    "            if item.is_file():\n",
+    "                shutil.copy2(item, trajectory_path)\n",
+    "            elif item.is_dir():\n",
+    "                shutil.copytree(item, trajectory_path / item.name, dirs_exist_ok=True)\n",
+    "    else:\n",
+    "        shutil.copy2(dataset, trajectory_path)\n",
     "    # Augment incoming task with CoT instructions\n",
     "    augmented_task = f\"\"\"\\\n",
     "    Here is the user query to address:\n",
@@ -96,7 +99,7 @@
     "# This image includes many bioinformatics and data science packages\n",
     "cfg.USE_DOCKER = False\n",
     "# This can be R or PYTHON in Docker or with a local kernel if you have R installed\n",
-    "LANGUAGE = NBLanguage.R\n",
+    "LANGUAGE = NBLanguage.PYTHON\n",
     "MAX_STEPS = 3\n",
     "MODEL_NAME = \"claude-3-7-sonnet-latest\""
    ]
@@ -115,9 +118,9 @@
     "logger = logging.getLogger(__name__)\n",
     "logger.info(\"Setting up data analysis environment\")\n",
     "\n",
-    "dataset_folder = Path(\"dataset\")\n",
+    "dataset = Path(\"datasets/brain_size_data.csv\")\n",
     "query = \"Analyze the dataset and give me an in depth analysis using pretty plots. I am particularly interested in crows.\"\n",
-    "environment = setup_data_analysis_env(query, dataset_folder, LANGUAGE)\n",
+    "environment = setup_data_analysis_env(query, dataset, LANGUAGE)\n",
     "\n",
     "agent = AgentConfig(\n",
     "    agent_type=\"ReActAgent\",\n",