diff --git a/docs/zero_to_hero_guide/03_Tool_Calling101.ipynb b/docs/zero_to_hero_guide/03_Tool_Calling101.ipynb
index f82482e01..6431c3e0a 100644
--- a/docs/zero_to_hero_guide/03_Tool_Calling101.ipynb
+++ b/docs/zero_to_hero_guide/03_Tool_Calling101.ipynb
@@ -19,51 +19,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 22,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Requirement already satisfied: llama-stack-client in ./.conda/envs/quick/lib/python3.13/site-packages (0.0.48)\n",
-      "Requirement already satisfied: anyio<5,>=3.5.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (4.6.2.post1)\n",
-      "Requirement already satisfied: distro<2,>=1.7.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (1.9.0)\n",
-      "Requirement already satisfied: httpx<1,>=0.23.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (0.27.2)\n",
-      "Requirement already satisfied: pydantic<3,>=1.9.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (2.9.2)\n",
-      "Requirement already satisfied: sniffio in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (1.3.1)\n",
-      "Requirement already satisfied: tabulate>=0.9.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (0.9.0)\n",
-      "Requirement already satisfied: typing-extensions<5,>=4.7 in ./.conda/envs/quick/lib/python3.13/site-packages (from llama-stack-client) (4.12.2)\n",
-      "Requirement already satisfied: idna>=2.8 in ./.conda/envs/quick/lib/python3.13/site-packages (from anyio<5,>=3.5.0->llama-stack-client) (3.10)\n",
-      "Requirement already satisfied: certifi in ./.conda/envs/quick/lib/python3.13/site-packages (from httpx<1,>=0.23.0->llama-stack-client) (2024.8.30)\n",
-      "Requirement already satisfied: httpcore==1.* in ./.conda/envs/quick/lib/python3.13/site-packages (from httpx<1,>=0.23.0->llama-stack-client) (1.0.6)\n",
-      "Requirement already satisfied: h11<0.15,>=0.13 in ./.conda/envs/quick/lib/python3.13/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->llama-stack-client) (0.14.0)\n",
-      "Requirement already satisfied: annotated-types>=0.6.0 in ./.conda/envs/quick/lib/python3.13/site-packages (from pydantic<3,>=1.9.0->llama-stack-client) (0.7.0)\n",
-      "Requirement already satisfied: pydantic-core==2.23.4 in ./.conda/envs/quick/lib/python3.13/site-packages (from pydantic<3,>=1.9.0->llama-stack-client) (2.23.4)\n"
-     ]
-    }
-   ],
-   "source": [
-    "!pip install llama-stack-client --upgrade"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "NameError",
-     "evalue": "name 'Agent' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[4], line 23\u001b[0m\n\u001b[1;32m     15\u001b[0m load_dotenv()\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# Helper function to create an agent with tools\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcreate_tool_agent\u001b[39m(\n\u001b[1;32m     19\u001b[0m     client: LlamaStackClient,\n\u001b[1;32m     20\u001b[0m     tools: List[Dict],\n\u001b[1;32m     21\u001b[0m     instructions: \u001b[38;5;28mstr\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou are a helpful assistant\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     22\u001b[0m     model: \u001b[38;5;28mstr\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLlama3.1-8B-Instruct\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m---> 23\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[43mAgent\u001b[49m:\n\u001b[1;32m     24\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Create an agent with specified tools.\"\"\"\u001b[39;00m\n\u001b[1;32m     25\u001b[0m     agent_config \u001b[38;5;241m=\u001b[39m AgentConfig(\n\u001b[1;32m     26\u001b[0m         model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m     27\u001b[0m         instructions\u001b[38;5;241m=\u001b[39minstructions,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     38\u001b[0m         enable_session_persistence\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     39\u001b[0m     )\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'Agent' is not defined"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import asyncio\n",
     "import os\n",
@@ -71,7 +29,7 @@
     "from dotenv import load_dotenv\n",
     "\n",
     "from llama_stack_client import LlamaStackClient\n",
-    "#from llama_stack_client.lib.agents.agent import Agent\n",
+    "from llama_stack_client.lib.agents.agent import Agent\n",
     "from llama_stack_client.lib.agents.event_logger import EventLogger\n",
     "from llama_stack_client.types.agent_create_params import (\n",
     "    AgentConfig,\n",
@@ -100,8 +58,6 @@
     "        tools=tools,\n",
     "        tool_choice=\"auto\",\n",
     "        tool_prompt_format=\"json\",\n",
-    "        input_shields=[\"llama_guard\"],\n",
-    "        output_shields=[\"llama_guard\"],\n",
     "        enable_session_persistence=True,\n",
     "    )\n",
     "\n",
@@ -121,16 +77,36 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Query: What are the latest developments in quantum computing?\n",
+      "--------------------------------------------------\n",
+      "\u001b[30m\u001b[0m\u001b[33minference> \u001b[0m\u001b[33mF\u001b[0m\u001b[33mIND\u001b[0m\u001b[33mINGS\u001b[0m\u001b[33m:\n",
+      "\u001b[0m\u001b[33mThe\u001b[0m\u001b[33m latest\u001b[0m\u001b[33m developments\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computing\u001b[0m\u001b[33m include\u001b[0m\u001b[33m advancements\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m processors\u001b[0m\u001b[33m,\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m algorithms\u001b[0m\u001b[33m,\u001b[0m\u001b[33m and\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m error\u001b[0m\u001b[33m correction\u001b[0m\u001b[33m.\u001b[0m\u001b[33m Researchers\u001b[0m\u001b[33m have\u001b[0m\u001b[33m made\u001b[0m\u001b[33m significant\u001b[0m\u001b[33m progress\u001b[0m\u001b[33m in\u001b[0m\u001b[33m developing\u001b[0m\u001b[33m more\u001b[0m\u001b[33m powerful\u001b[0m\u001b[33m and\u001b[0m\u001b[33m reliable\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computers\u001b[0m\u001b[33m,\u001b[0m\u001b[33m with\u001b[0m\u001b[33m some\u001b[0m\u001b[33m companies\u001b[0m\u001b[33m already\u001b[0m\u001b[33m showcasing\u001b[0m\u001b[33m \u001b[0m\u001b[33m100\u001b[0m\u001b[33m-q\u001b[0m\u001b[33mubit\u001b[0m\u001b[33m and\u001b[0m\u001b[33m \u001b[0m\u001b[33m127\u001b[0m\u001b[33m-q\u001b[0m\u001b[33mubit\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m processors\u001b[0m\u001b[33m (\u001b[0m\u001b[33mIBM\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m;\u001b[0m\u001b[33m Google\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\u001b[0m\u001b[33m These\u001b[0m\u001b[33m advancements\u001b[0m\u001b[33m have\u001b[0m\u001b[33m led\u001b[0m\u001b[33m to\u001b[0m\u001b[33m breakthrough\u001b[0m\u001b[33ms\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m simulations\u001b[0m\u001b[33m,\u001b[0m\u001b[33m machine\u001b[0m\u001b[33m learning\u001b[0m\u001b[33m,\u001b[0m\u001b[33m and\u001b[0m\u001b[33m optimization\u001b[0m\u001b[33m problems\u001b[0m\u001b[33m (\u001b[0m\u001b[33mB\u001b[0m\u001b[33mhart\u001b[0m\u001b[33mi\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m;\u001b[0m\u001b[33m Zhang\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\u001b[0m\u001b[33m Additionally\u001b[0m\u001b[33m,\u001b[0m\u001b[33m there\u001b[0m\u001b[33m have\u001b[0m\u001b[33m been\u001b[0m\u001b[33m significant\u001b[0m\u001b[33m improvements\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m error\u001b[0m\u001b[33m correction\u001b[0m\u001b[33m,\u001b[0m\u001b[33m which\u001b[0m\u001b[33m is\u001b[0m\u001b[33m essential\u001b[0m\u001b[33m for\u001b[0m\u001b[33m large\u001b[0m\u001b[33m-scale\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computing\u001b[0m\u001b[33m (\u001b[0m\u001b[33mG\u001b[0m\u001b[33mottes\u001b[0m\u001b[33mman\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\n",
+      "\n",
+      "\u001b[0m\u001b[33mS\u001b[0m\u001b[33mOURCES\u001b[0m\u001b[33m:\n",
+      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m IBM\u001b[0m\u001b[33m Quantum\u001b[0m\u001b[33m:\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m Process\u001b[0m\u001b[33mors\u001b[0m\u001b[33m\"\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mwww\u001b[0m\u001b[33m.ibm\u001b[0m\u001b[33m.com\u001b[0m\u001b[33m/\u001b[0m\u001b[33mquant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m/com\u001b[0m\u001b[33mputer\u001b[0m\u001b[33m/)\n",
+      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m Google\u001b[0m\u001b[33m Quantum\u001b[0m\u001b[33m AI\u001b[0m\u001b[33m Lab\u001b[0m\u001b[33m:\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m Process\u001b[0m\u001b[33mors\u001b[0m\u001b[33m\"\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mquant\u001b[0m\u001b[33mum\u001b[0m\u001b[33mai\u001b[0m\u001b[33m.google\u001b[0m\u001b[33m/al\u001b[0m\u001b[33mphabet\u001b[0m\u001b[33m/sub\u001b[0m\u001b[33m-page\u001b[0m\u001b[33m-\u001b[0m\u001b[33m1\u001b[0m\u001b[33m/)\n",
+      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m Bh\u001b[0m\u001b[33marti\u001b[0m\u001b[33m,\u001b[0m\u001b[33m K\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m Computing\u001b[0m\u001b[33m:\u001b[0m\u001b[33m A\u001b[0m\u001b[33m Review\u001b[0m\u001b[33m of\u001b[0m\u001b[33m Recent\u001b[0m\u001b[33m Advances\u001b[0m\u001b[33m.\"\u001b[0m\u001b[33m Journal\u001b[0m\u001b[33m of\u001b[0m\u001b[33m Physics\u001b[0m\u001b[33m:\u001b[0m\u001b[33m Conference\u001b[0m\u001b[33m Series\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m218\u001b[0m\u001b[33m5\u001b[0m\u001b[33m(\u001b[0m\u001b[33m1\u001b[0m\u001b[33m),\u001b[0m\u001b[33m \u001b[0m\u001b[33m012\u001b[0m\u001b[33m001\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mi\u001b[0m\u001b[33mop\u001b[0m\u001b[33mscience\u001b[0m\u001b[33m.i\u001b[0m\u001b[33mop\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/article\u001b[0m\u001b[33m/\u001b[0m\u001b[33m10\u001b[0m\u001b[33m.\u001b[0m\u001b[33m108\u001b[0m\u001b[33m8\u001b[0m\u001b[33m/\u001b[0m\u001b[33m174\u001b[0m\u001b[33m2\u001b[0m\u001b[33m-\u001b[0m\u001b[33m659\u001b[0m\u001b[33m6\u001b[0m\u001b[33m/\u001b[0m\u001b[33m218\u001b[0m\u001b[33m5\u001b[0m\u001b[33m/\u001b[0m\u001b[33m1\u001b[0m\u001b[33m/\u001b[0m\u001b[33m012\u001b[0m\u001b[33m001\u001b[0m\u001b[33m)\n",
+      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m Zhang\u001b[0m\u001b[33m,\u001b[0m\u001b[33m Y\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m Algorithms\u001b[0m\u001b[33m for\u001b[0m\u001b[33m Machine\u001b[0m\u001b[33m Learning\u001b[0m\u001b[33m.\"\u001b[0m\u001b[33m Journal\u001b[0m\u001b[33m of\u001b[0m\u001b[33m Machine\u001b[0m\u001b[33m Learning\u001b[0m\u001b[33m Research\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m23\u001b[0m\u001b[33m,\u001b[0m\u001b[33m \u001b[0m\u001b[33m1\u001b[0m\u001b[33m-\u001b[0m\u001b[33m36\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mj\u001b[0m\u001b[33mml\u001b[0m\u001b[33mr\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/p\u001b[0m\u001b[33mapers\u001b[0m\u001b[33m/v\u001b[0m\u001b[33m23\u001b[0m\u001b[33m/\u001b[0m\u001b[33m20\u001b[0m\u001b[33m-\u001b[0m\u001b[33m065\u001b[0m\u001b[33m.html\u001b[0m\u001b[33m)\n",
+      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m G\u001b[0m\u001b[33mottes\u001b[0m\u001b[33mman\u001b[0m\u001b[33m,\u001b[0m\u001b[33m D\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33m202\u001b[0m\u001b[33m2\u001b[0m\u001b[33m).\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m Error\u001b[0m\u001b[33m Correction\u001b[0m\u001b[33m.\"\u001b[0m\u001b[33m In\u001b[0m\u001b[33m Encyclopedia\u001b[0m\u001b[33m of\u001b[0m\u001b[33m Complexity\u001b[0m\u001b[33m and\u001b[0m\u001b[33m Systems\u001b[0m\u001b[33m Science\u001b[0m\u001b[33m (\u001b[0m\u001b[33mpp\u001b[0m\u001b[33m.\u001b[0m\u001b[33m \u001b[0m\u001b[33m1\u001b[0m\u001b[33m-\u001b[0m\u001b[33m13\u001b[0m\u001b[33m).\u001b[0m\u001b[33m Springer\u001b[0m\u001b[33m,\u001b[0m\u001b[33m New\u001b[0m\u001b[33m York\u001b[0m\u001b[33m,\u001b[0m\u001b[33m NY\u001b[0m\u001b[33m.\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mlink\u001b[0m\u001b[33m.spring\u001b[0m\u001b[33mer\u001b[0m\u001b[33m.com\u001b[0m\u001b[33m/reference\u001b[0m\u001b[33mwork\u001b[0m\u001b[33mentry\u001b[0m\u001b[33m/\u001b[0m\u001b[33m10\u001b[0m\u001b[33m.\u001b[0m\u001b[33m100\u001b[0m\u001b[33m7\u001b[0m\u001b[33m/\u001b[0m\u001b[33m978\u001b[0m\u001b[33m-\u001b[0m\u001b[33m0\u001b[0m\u001b[33m-\u001b[0m\u001b[33m387\u001b[0m\u001b[33m-\u001b[0m\u001b[33m758\u001b[0m\u001b[33m88\u001b[0m\u001b[33m-\u001b[0m\u001b[33m6\u001b[0m\u001b[33m_\u001b[0m\u001b[33m447\u001b[0m\u001b[33m)\u001b[0m\u001b[97m\u001b[0m\n",
+      "\u001b[30m\u001b[0m"
+     ]
+    }
+   ],
    "source": [
     "async def create_search_agent(client: LlamaStackClient) -> Agent:\n",
     "    \"\"\"Create an agent with Brave Search capability.\"\"\"\n",
     "    search_tool = AgentConfigToolSearchToolDefinition(\n",
     "        type=\"brave_search\",\n",
     "        engine=\"brave\",\n",
-    "        api_key=os.getenv(\"BRAVE_SEARCH_API_KEY\"),\n",
+    "        api_key=\"dummy_value\"#os.getenv(\"BRAVE_SEARCH_API_KEY\"),\n",
     "    )\n",
     "\n",
     "    return await create_tool_agent(\n",
@@ -151,7 +127,7 @@
     "\n",
     "# Example usage\n",
     "async def search_example():\n",
-    "    client = LlamaStackClient(base_url=\"http://localhost:8000\")\n",
+    "    client = LlamaStackClient(base_url=\"http://localhost:5001\")\n",
     "    agent = await create_search_agent(client)\n",
     "\n",
     "    # Create a session\n",
@@ -160,7 +136,7 @@
     "    # Example queries\n",
     "    queries = [\n",
     "        \"What are the latest developments in quantum computing?\",\n",
-    "        \"Who won the most recent Super Bowl?\",\n",
+    "        #\"Who won the most recent Super Bowl?\",\n",
     "    ]\n",
     "\n",
     "    for query in queries:\n",
@@ -190,30 +166,68 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 27,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Query: What's the weather like in San Francisco?\n",
+      "--------------------------------------------------\n",
+      "\u001b[30m\u001b[0m\u001b[33minference> \u001b[0m\u001b[33m{\n",
+      "\u001b[0m\u001b[33m   \u001b[0m\u001b[33m \"\u001b[0m\u001b[33mtype\u001b[0m\u001b[33m\":\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mfunction\u001b[0m\u001b[33m\",\n",
+      "\u001b[0m\u001b[33m   \u001b[0m\u001b[33m \"\u001b[0m\u001b[33mname\u001b[0m\u001b[33m\":\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mget\u001b[0m\u001b[33m_weather\u001b[0m\u001b[33m\",\n",
+      "\u001b[0m\u001b[33m   \u001b[0m\u001b[33m \"\u001b[0m\u001b[33mparameters\u001b[0m\u001b[33m\":\u001b[0m\u001b[33m {\n",
+      "\u001b[0m\u001b[33m       \u001b[0m\u001b[33m \"\u001b[0m\u001b[33mlocation\u001b[0m\u001b[33m\":\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mSan\u001b[0m\u001b[33m Francisco\u001b[0m\u001b[33m\"\n",
+      "\u001b[0m\u001b[33m   \u001b[0m\u001b[33m }\n",
+      "\u001b[0m\u001b[33m}\u001b[0m\u001b[97m\u001b[0m\n"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'WeatherTool' object has no attribute 'run'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[27], line 113\u001b[0m\n\u001b[1;32m    110\u001b[0m nest_asyncio\u001b[38;5;241m.\u001b[39mapply()\n\u001b[1;32m    112\u001b[0m \u001b[38;5;66;03m# Run the example\u001b[39;00m\n\u001b[0;32m--> 113\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m weather_example()\n",
+      "Cell \u001b[0;32mIn[27], line 105\u001b[0m, in \u001b[0;36mweather_example\u001b[0;34m()\u001b[0m\n\u001b[1;32m     98\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m-\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m50\u001b[39m)\n\u001b[1;32m    100\u001b[0m response \u001b[38;5;241m=\u001b[39m agent\u001b[38;5;241m.\u001b[39mcreate_turn(\n\u001b[1;32m    101\u001b[0m     messages\u001b[38;5;241m=\u001b[39m[{\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrole\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124muser\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcontent\u001b[39m\u001b[38;5;124m\"\u001b[39m: query}],\n\u001b[1;32m    102\u001b[0m     session_id\u001b[38;5;241m=\u001b[39msession_id,\n\u001b[1;32m    103\u001b[0m )\n\u001b[0;32m--> 105\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m log \u001b[38;5;129;01min\u001b[39;00m EventLogger()\u001b[38;5;241m.\u001b[39mlog(response):\n\u001b[1;32m    106\u001b[0m     log\u001b[38;5;241m.\u001b[39mprint()\n",
+      "File \u001b[0;32m~/new_task/llama-stack-client-python/src/llama_stack_client/lib/agents/event_logger.py:55\u001b[0m, in \u001b[0;36mEventLogger.log\u001b[0;34m(self, event_generator)\u001b[0m\n\u001b[1;32m     52\u001b[0m previous_event_type \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     53\u001b[0m previous_step_type \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 55\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m chunk \u001b[38;5;129;01min\u001b[39;00m event_generator:\n\u001b[1;32m     56\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(chunk, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mevent\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m     57\u001b[0m         \u001b[38;5;66;03m# Need to check for custom tool first\u001b[39;00m\n\u001b[1;32m     58\u001b[0m         \u001b[38;5;66;03m# since it does not produce event but instead\u001b[39;00m\n\u001b[1;32m     59\u001b[0m         \u001b[38;5;66;03m# a Message\u001b[39;00m\n\u001b[1;32m     60\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(chunk, ToolResponseMessage):\n",
+      "File \u001b[0;32m~/new_task/llama-stack-client-python/src/llama_stack_client/lib/agents/agent.py:76\u001b[0m, in \u001b[0;36mAgent.create_turn\u001b[0;34m(self, messages, attachments, session_id)\u001b[0m\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     75\u001b[0m     tool \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcustom_tools[tool_call\u001b[38;5;241m.\u001b[39mtool_name]\n\u001b[0;32m---> 76\u001b[0m     result_messages \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mexecute_custom_tool(tool, message)\n\u001b[1;32m     77\u001b[0m     next_message \u001b[38;5;241m=\u001b[39m result_messages[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28;01myield\u001b[39;00m next_message\n",
+      "File \u001b[0;32m~/new_task/llama-stack-client-python/src/llama_stack_client/lib/agents/agent.py:84\u001b[0m, in \u001b[0;36mAgent.execute_custom_tool\u001b[0;34m(self, tool, message)\u001b[0m\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mexecute_custom_tool\u001b[39m(\n\u001b[1;32m     82\u001b[0m     \u001b[38;5;28mself\u001b[39m, tool: CustomTool, message: Union[UserMessage, ToolResponseMessage]\n\u001b[1;32m     83\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m List[Union[UserMessage, ToolResponseMessage]]:\n\u001b[0;32m---> 84\u001b[0m     result_messages \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[43mtool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m([message])\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result_messages\n",
+      "\u001b[0;31mAttributeError\u001b[0m: 'WeatherTool' object has no attribute 'run'"
+     ]
+    }
+   ],
    "source": [
-    "from typing import TypedDict, Optional\n",
+    "from typing import TypedDict, Optional, Dict, Any\n",
     "from datetime import datetime\n",
-    "\n",
-    "# Define tool types\n",
-    "class WeatherInput(TypedDict):\n",
-    "    location: str\n",
-    "    date: Optional[str]\n",
-    "\n",
-    "class WeatherOutput(TypedDict):\n",
-    "    temperature: float\n",
-    "    conditions: str\n",
-    "    humidity: float\n",
-    "\n",
     "class WeatherTool:\n",
     "    \"\"\"Example custom tool for weather information.\"\"\"\n",
-    "\n",
-    "    def __init__(self, api_key: Optional[str] = None):\n",
-    "        self.api_key = api_key\n",
-    "\n",
-    "    async def get_weather(self, location: str, date: Optional[str] = None) -> WeatherOutput:\n",
+    "    \n",
+    "    def get_name(self) -> str:\n",
+    "        return \"get_weather\"\n",
+    "    \n",
+    "    def get_description(self) -> str:\n",
+    "        return \"Get weather information for a location\"\n",
+    "        \n",
+    "    def get_params_definition(self) -> Dict[str, ToolParamDefinitionParam]:\n",
+    "        return {\n",
+    "            \"location\": ToolParamDefinitionParam(\n",
+    "                param_type=\"str\",\n",
+    "                description=\"City or location name\",\n",
+    "                required=True\n",
+    "            ),\n",
+    "            \"date\": ToolParamDefinitionParam(\n",
+    "                param_type=\"str\",\n",
+    "                description=\"Optional date (YYYY-MM-DD)\",\n",
+    "                required=False\n",
+    "            )\n",
+    "        }\n",
+    "    \n",
+    "    async def run_impl(self, location: str, date: Optional[str] = None) -> Dict[str, Any]:\n",
     "        \"\"\"Simulate getting weather data (replace with actual API call).\"\"\"\n",
     "        # Mock implementation\n",
     "        return {\n",
@@ -222,76 +236,93 @@
     "            \"humidity\": 65.0\n",
     "        }\n",
     "\n",
-    "    async def __call__(self, input_data: WeatherInput) -> WeatherOutput:\n",
-    "        \"\"\"Make the tool callable with structured input.\"\"\"\n",
-    "        return await self.get_weather(\n",
-    "            location=input_data[\"location\"],\n",
-    "            date=input_data.get(\"date\")\n",
-    "        )\n",
-    "\n",
     "async def create_weather_agent(client: LlamaStackClient) -> Agent:\n",
     "    \"\"\"Create an agent with weather tool capability.\"\"\"\n",
-    "    weather_tool = {\n",
-    "        \"type\": \"function\",\n",
-    "        \"function\": {\n",
-    "            \"name\": \"get_weather\",\n",
-    "            \"description\": \"Get weather information for a location\",\n",
-    "            \"parameters\": {\n",
-    "                \"type\": \"object\",\n",
-    "                \"properties\": {\n",
-    "                    \"location\": {\n",
-    "                        \"type\": \"string\",\n",
-    "                        \"description\": \"City or location name\"\n",
-    "                    },\n",
-    "                    \"date\": {\n",
-    "                        \"type\": \"string\",\n",
-    "                        \"description\": \"Optional date (YYYY-MM-DD)\",\n",
-    "                        \"format\": \"date\"\n",
-    "                    }\n",
-    "                },\n",
-    "                \"required\": [\"location\"]\n",
-    "            }\n",
-    "        },\n",
-    "        \"implementation\": WeatherTool()\n",
-    "    }\n",
-    "\n",
-    "    return await create_tool_agent(\n",
-    "        client=client,\n",
-    "        tools=[weather_tool],\n",
+    "    agent_config = AgentConfig(\n",
+    "        model=\"Llama3.1-8B-Instruct\",\n",
     "        instructions=\"\"\"\n",
     "        You are a weather assistant that can provide weather information.\n",
     "        Always specify the location clearly in your responses.\n",
     "        Include both temperature and conditions in your summaries.\n",
-    "        \"\"\"\n",
+    "        \"\"\",\n",
+    "        sampling_params={\n",
+    "            \"strategy\": \"greedy\",\n",
+    "            \"temperature\": 1.0,\n",
+    "            \"top_p\": 0.9,\n",
+    "        },\n",
+    "        tools=[\n",
+    "            {\n",
+    "                \"function_name\": \"get_weather\",\n",
+    "                \"description\": \"Get weather information for a location\",\n",
+    "                \"parameters\": {\n",
+    "                    \"location\": {\n",
+    "                        \"param_type\": \"str\",\n",
+    "                        \"description\": \"City or location name\",\n",
+    "                        \"required\": True,\n",
+    "                    },\n",
+    "                    \"date\": {\n",
+    "                        \"param_type\": \"str\",\n",
+    "                        \"description\": \"Optional date (YYYY-MM-DD)\",\n",
+    "                        \"required\": False,\n",
+    "                    },\n",
+    "                },\n",
+    "                \"type\": \"function_call\",\n",
+    "            }\n",
+    "        ],\n",
+    "        tool_choice=\"auto\",\n",
+    "        tool_prompt_format=\"json\",\n",
+    "        input_shields=[],\n",
+    "        output_shields=[],\n",
+    "        enable_session_persistence=True\n",
     "    )\n",
+    "    \n",
+    "    # Create the agent with the tool\n",
+    "    weather_tool = WeatherTool()\n",
+    "    agent = Agent(\n",
+    "        client=client,\n",
+    "        agent_config=agent_config,\n",
+    "        custom_tools=[weather_tool]\n",
+    "    )\n",
+    "    \n",
+    "    return agent\n",
     "\n",
     "# Example usage\n",
     "async def weather_example():\n",
-    "    client = LlamaStackClient(base_url=\"http://localhost:8000\")\n",
+    "    client = LlamaStackClient(base_url=\"http://localhost:5001\")\n",
     "    agent = await create_weather_agent(client)\n",
-    "\n",
     "    session_id = agent.create_session(\"weather-session\")\n",
-    "\n",
+    "    \n",
     "    queries = [\n",
     "        \"What's the weather like in San Francisco?\",\n",
     "        \"Tell me the weather in Tokyo tomorrow\",\n",
     "    ]\n",
-    "\n",
+    "    \n",
     "    for query in queries:\n",
     "        print(f\"\\nQuery: {query}\")\n",
     "        print(\"-\" * 50)\n",
-    "\n",
+    "        \n",
     "        response = agent.create_turn(\n",
     "            messages=[{\"role\": \"user\", \"content\": query}],\n",
     "            session_id=session_id,\n",
     "        )\n",
-    "\n",
+    "        \n",
     "        async for log in EventLogger().log(response):\n",
     "            log.print()\n",
     "\n",
+    "# For Jupyter notebooks\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()\n",
+    "\n",
     "# Run the example\n",
     "await weather_example()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {