{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tool Calling\n",
    "\n",
    "Before you begin, please ensure Llama Stack is installed and set up by following the [Getting Started Guide](https://llama-stack.readthedocs.io/en/latest/getting_started/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we'll explore how to enhance your applications with tool calling capabilities. We'll cover:\n",
    "1. Setting up and using the Brave Search API\n",
    "2. Creating custom tools\n",
    "3. Configuring tool prompts and safety settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up your connection parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOST = \"localhost\"  # Replace with your host\n",
    "PORT = 5000        # Replace with your port"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import os\n",
    "from typing import Dict, List, Optional\n",
    "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.event_logger import EventLogger\n",
    "from llama_stack_client.types.agent_create_params import (\n",
    "    AgentConfig,\n",
    "    AgentConfigToolSearchToolDefinition,\n",
    ")\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Helper function to create an agent with tools\n",
    "async def create_tool_agent(\n",
    "    client: LlamaStackClient,\n",
    "    tools: List[Dict],\n",
    "    instructions: str = \"You are a helpful assistant\",\n",
    "    model: str = \"Llama3.2-11B-Vision-Instruct\",\n",
    ") -> Agent:\n",
    "    \"\"\"Create an agent with specified tools.\"\"\"\n",
    "    print(\"Using the following model: \", model)\n",
    "    agent_config = AgentConfig(\n",
    "        model=model,\n",
    "        instructions=instructions,\n",
    "        sampling_params={\n",
    "            \"strategy\": \"greedy\",\n",
    "            \"temperature\": 1.0,\n",
    "            \"top_p\": 0.9,\n",
    "        },\n",
    "        tools=tools,\n",
    "        tool_choice=\"auto\",\n",
    "        tool_prompt_format=\"json\",\n",
    "        enable_session_persistence=True,\n",
    "    )\n",
    "\n",
    "    return Agent(client, agent_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, create a `.env` file in your notebook directory with your Brave Search API key:\n",
    "\n",
    "```\n",
    "BRAVE_SEARCH_API_KEY=your_key_here\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using the following model:  Llama3.2-11B-Vision-Instruct\n",
      "\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 the\u001b[0m\u001b[33m introduction\u001b[0m\u001b[33m of\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m processors\u001b[0m\u001b[33m with\u001b[0m\u001b[33m increasing\u001b[0m\u001b[33m numbers\u001b[0m\u001b[33m of\u001b[0m\u001b[33m q\u001b[0m\u001b[33mubits\u001b[0m\u001b[33m,\u001b[0m\u001b[33m advancements\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 and\u001b[0m\u001b[33m the\u001b[0m\u001b[33m development\u001b[0m\u001b[33m of\u001b[0m\u001b[33m new\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m algorithms\u001b[0m\u001b[33m.\u001b[0m\u001b[33m Companies\u001b[0m\u001b[33m such\u001b[0m\u001b[33m as\u001b[0m\u001b[33m IBM\u001b[0m\u001b[33m,\u001b[0m\u001b[33m Google\u001b[0m\u001b[33m,\u001b[0m\u001b[33m and\u001b[0m\u001b[33m Microsoft\u001b[0m\u001b[33m are\u001b[0m\u001b[33m leading\u001b[0m\u001b[33m the\u001b[0m\u001b[33m charge\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computing\u001b[0m\u001b[33m research\u001b[0m\u001b[33m and\u001b[0m\u001b[33m development\u001b[0m\u001b[33m.\u001b[0m\u001b[33m For\u001b[0m\u001b[33m example\u001b[0m\u001b[33m,\u001b[0m\u001b[33m IBM\u001b[0m\u001b[33m's\u001b[0m\u001b[33m Quantum\u001b[0m\u001b[33m Experience\u001b[0m\u001b[33m cloud\u001b[0m\u001b[33m-based\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computer\u001b[0m\u001b[33m has\u001b[0m\u001b[33m \u001b[0m\u001b[33m127\u001b[0m\u001b[33m q\u001b[0m\u001b[33mubits\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[33mput\u001b[0m\u001b[33ming\u001b[0m\u001b[33m/\u001b[0m\u001b[33m),\u001b[0m\u001b[33m while\u001b[0m\u001b[33m Google\u001b[0m\u001b[33m's\u001b[0m\u001b[33m S\u001b[0m\u001b[33myc\u001b[0m\u001b[33mam\u001b[0m\u001b[33more\u001b[0m\u001b[33m processor\u001b[0m\u001b[33m has\u001b[0m\u001b[33m \u001b[0m\u001b[33m53\u001b[0m\u001b[33m q\u001b[0m\u001b[33mubits\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33ms\u001b[0m\u001b[33mcai\u001b[0m\u001b[33m.st\u001b[0m\u001b[33manford\u001b[0m\u001b[33m.edu\u001b[0m\u001b[33m/~\u001b[0m\u001b[33mmw\u001b[0m\u001b[33mj\u001b[0m\u001b[33m/s\u001b[0m\u001b[33myc\u001b[0m\u001b[33mam\u001b[0m\u001b[33more\u001b[0m\u001b[33m.html\u001b[0m\u001b[33m).\n",
      "\n",
      "\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m error\u001b[0m\u001b[33m correction\u001b[0m\u001b[33m is\u001b[0m\u001b[33m also\u001b[0m\u001b[33m a\u001b[0m\u001b[33m major\u001b[0m\u001b[33m area\u001b[0m\u001b[33m of\u001b[0m\u001b[33m research\u001b[0m\u001b[33m,\u001b[0m\u001b[33m with\u001b[0m\u001b[33m scientists\u001b[0m\u001b[33m exploring\u001b[0m\u001b[33m various\u001b[0m\u001b[33m methods\u001b[0m\u001b[33m to\u001b[0m\u001b[33m mitigate\u001b[0m\u001b[33m the\u001b[0m\u001b[33m effects\u001b[0m\u001b[33m of\u001b[0m\u001b[33m deco\u001b[0m\u001b[33mherence\u001b[0m\u001b[33m and\u001b[0m\u001b[33m errors\u001b[0m\u001b[33m in\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computations\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mar\u001b[0m\u001b[33mxiv\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/\u001b[0m\u001b[33mabs\u001b[0m\u001b[33m/\u001b[0m\u001b[33m210\u001b[0m\u001b[33m4\u001b[0m\u001b[33m.\u001b[0m\u001b[33m045\u001b[0m\u001b[33m93\u001b[0m\u001b[33m).\u001b[0m\u001b[33m New\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m algorithms\u001b[0m\u001b[33m,\u001b[0m\u001b[33m such\u001b[0m\u001b[33m as\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m simulation\u001b[0m\u001b[33m and\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m machine\u001b[0m\u001b[33m learning\u001b[0m\u001b[33m,\u001b[0m\u001b[33m are\u001b[0m\u001b[33m also\u001b[0m\u001b[33m being\u001b[0m\u001b[33m developed\u001b[0m\u001b[33m to\u001b[0m\u001b[33m take\u001b[0m\u001b[33m advantage\u001b[0m\u001b[33m of\u001b[0m\u001b[33m the\u001b[0m\u001b[33m unique\u001b[0m\u001b[33m properties\u001b[0m\u001b[33m of\u001b[0m\u001b[33m quantum\u001b[0m\u001b[33m computing\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mar\u001b[0m\u001b[33mxiv\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/\u001b[0m\u001b[33mabs\u001b[0m\u001b[33m/\u001b[0m\u001b[33m210\u001b[0m\u001b[33m2\u001b[0m\u001b[33m.\u001b[0m\u001b[33m081\u001b[0m\u001b[33m83\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 Experience\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[33mput\u001b[0m\u001b[33ming\u001b[0m\u001b[33m/)\n",
      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m Google\u001b[0m\u001b[33m's\u001b[0m\u001b[33m S\u001b[0m\u001b[33myc\u001b[0m\u001b[33mam\u001b[0m\u001b[33more\u001b[0m\u001b[33m processor\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33ms\u001b[0m\u001b[33mcai\u001b[0m\u001b[33m.st\u001b[0m\u001b[33manford\u001b[0m\u001b[33m.edu\u001b[0m\u001b[33m/~\u001b[0m\u001b[33mmw\u001b[0m\u001b[33mj\u001b[0m\u001b[33m/s\u001b[0m\u001b[33myc\u001b[0m\u001b[33mam\u001b[0m\u001b[33more\u001b[0m\u001b[33m.html\u001b[0m\u001b[33m)\n",
      "\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 with\u001b[0m\u001b[33m encoded\u001b[0m\u001b[33m cat\u001b[0m\u001b[33m codes\u001b[0m\u001b[33m\"\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mar\u001b[0m\u001b[33mxiv\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/\u001b[0m\u001b[33mabs\u001b[0m\u001b[33m/\u001b[0m\u001b[33m210\u001b[0m\u001b[33m4\u001b[0m\u001b[33m.\u001b[0m\u001b[33m045\u001b[0m\u001b[33m93\u001b[0m\u001b[33m)\n",
      "\u001b[0m\u001b[33m-\u001b[0m\u001b[33m \"\u001b[0m\u001b[33mQuant\u001b[0m\u001b[33mum\u001b[0m\u001b[33m simulation\u001b[0m\u001b[33m of\u001b[0m\u001b[33m many\u001b[0m\u001b[33m-body\u001b[0m\u001b[33m systems\u001b[0m\u001b[33m\"\u001b[0m\u001b[33m (\u001b[0m\u001b[33mhttps\u001b[0m\u001b[33m://\u001b[0m\u001b[33mar\u001b[0m\u001b[33mxiv\u001b[0m\u001b[33m.org\u001b[0m\u001b[33m/\u001b[0m\u001b[33mabs\u001b[0m\u001b[33m/\u001b[0m\u001b[33m210\u001b[0m\u001b[33m2\u001b[0m\u001b[33m.\u001b[0m\u001b[33m081\u001b[0m\u001b[33m83\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=\"dummy_value\"#os.getenv(\"BRAVE_SEARCH_API_KEY\"),\n",
    "    )\n",
    "\n",
    "    models_response = client.models.list()\n",
    "    for model in models_response:\n",
    "        if model.identifier.endswith(\"Instruct\"):\n",
    "            model_name = model.llama_model\n",
    "\n",
    "\n",
    "    return await create_tool_agent(\n",
    "        client=client,\n",
    "        tools=[search_tool],\n",
    "        model = model_name,\n",
    "        instructions=\"\"\"\n",
    "        You are a research assistant that can search the web.\n",
    "        Always cite your sources with URLs when providing information.\n",
    "        Format your responses as:\n",
    "\n",
    "        FINDINGS:\n",
    "        [Your summary here]\n",
    "\n",
    "        SOURCES:\n",
    "        - [Source title](URL)\n",
    "        \"\"\"\n",
    "    )\n",
    "\n",
    "# Example usage\n",
    "async def search_example():\n",
    "    client = LlamaStackClient(base_url=f\"http://{HOST}:{PORT}\")\n",
    "    agent = await create_search_agent(client)\n",
    "\n",
    "    # Create a session\n",
    "    session_id = agent.create_session(\"search-session\")\n",
    "\n",
    "    # Example queries\n",
    "    queries = [\n",
    "        \"What are the latest developments in quantum computing?\",\n",
    "        #\"Who won the most recent Super Bowl?\",\n",
    "    ]\n",
    "\n",
    "    for query in queries:\n",
    "        print(f\"\\nQuery: {query}\")\n",
    "        print(\"-\" * 50)\n",
    "\n",
    "        response = agent.create_turn(\n",
    "            messages=[{\"role\": \"user\", \"content\": query}],\n",
    "            session_id=session_id,\n",
    "        )\n",
    "\n",
    "        async for log in EventLogger().log(response):\n",
    "            log.print()\n",
    "\n",
    "# Run the example (in Jupyter, use asyncio.run())\n",
    "await search_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Custom Tool Creation\n",
    "\n",
    "Let's create a custom weather tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "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",
      "\u001b[32mCustomTool> {\"temperature\": 72.5, \"conditions\": \"partly cloudy\", \"humidity\": 65.0}\u001b[0m\n",
      "\n",
      "Query: Tell me the weather in Tokyo tomorrow\n",
      "--------------------------------------------------\n",
      "\u001b[30m\u001b[0m\u001b[33minference> \u001b[0m\u001b[36m\u001b[0m\u001b[36m{\"\u001b[0m\u001b[36mtype\u001b[0m\u001b[36m\":\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mfunction\u001b[0m\u001b[36m\",\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mname\u001b[0m\u001b[36m\":\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mget\u001b[0m\u001b[36m_weather\u001b[0m\u001b[36m\",\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mparameters\u001b[0m\u001b[36m\":\u001b[0m\u001b[36m {\"\u001b[0m\u001b[36mlocation\u001b[0m\u001b[36m\":\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mTok\u001b[0m\u001b[36myo\u001b[0m\u001b[36m\",\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mdate\u001b[0m\u001b[36m\":\u001b[0m\u001b[36m \"\u001b[0m\u001b[36mtom\u001b[0m\u001b[36morrow\u001b[0m\u001b[36m\"}}\u001b[0m\u001b[97m\u001b[0m\n",
      "\u001b[32mCustomTool> {\"temperature\": 90.1, \"conditions\": \"sunny\", \"humidity\": 40.0}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from typing import TypedDict, Optional, Dict, Any\n",
    "from datetime import datetime\n",
    "import json\n",
    "from llama_stack_client.types.tool_param_definition_param import ToolParamDefinitionParam\n",
    "from llama_stack_client.types import CompletionMessage,ToolResponseMessage\n",
    "from llama_stack_client.lib.agents.custom_tool import CustomTool\n",
    "\n",
    "class WeatherTool(CustomTool):\n",
    "    \"\"\"Example custom tool for weather information.\"\"\"\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",
    "    async def run(self, messages: List[CompletionMessage]) -> List[ToolResponseMessage]:\n",
    "        assert len(messages) == 1, \"Expected single message\"\n",
    "\n",
    "        message = messages[0]\n",
    "\n",
    "        tool_call = message.tool_calls[0]\n",
    "        # location = tool_call.arguments.get(\"location\", None)\n",
    "        # date = tool_call.arguments.get(\"date\", None)\n",
    "        try:\n",
    "            response = await self.run_impl(**tool_call.arguments)\n",
    "            response_str = json.dumps(response, ensure_ascii=False)\n",
    "        except Exception as e:\n",
    "            response_str = f\"Error when running tool: {e}\"\n",
    "\n",
    "        message = ToolResponseMessage(\n",
    "            call_id=tool_call.call_id,\n",
    "            tool_name=tool_call.tool_name,\n",
    "            content=response_str,\n",
    "            role=\"ipython\",\n",
    "        )\n",
    "        return [message]\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",
    "        if date:\n",
    "            return {\n",
    "            \"temperature\": 90.1,\n",
    "            \"conditions\": \"sunny\",\n",
    "            \"humidity\": 40.0\n",
    "        }\n",
    "        return {\n",
    "            \"temperature\": 72.5,\n",
    "            \"conditions\": \"partly cloudy\",\n",
    "            \"humidity\": 65.0\n",
    "        }\n",
    "\n",
    "\n",
    "async def create_weather_agent(client: LlamaStackClient) -> Agent:\n",
    "    \"\"\"Create an agent with weather tool capability.\"\"\"\n",
    "    models_response = client.models.list()\n",
    "    for model in models_response:\n",
    "        if model.identifier.endswith(\"Instruct\"):\n",
    "            model_name = model.llama_model\n",
    "    agent_config = AgentConfig(\n",
    "        model=model_name,\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",
    "        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=f\"http://{HOST}:{PORT}\")\n",
    "    agent = await create_weather_agent(client)\n",
    "    session_id = agent.create_session(\"weather-session\")\n",
    "\n",
    "    queries = [\n",
    "        \"What's the weather like in San Francisco?\",\n",
    "        \"Tell me the weather in Tokyo tomorrow\",\n",
    "    ]\n",
    "\n",
    "    for query in queries:\n",
    "        print(f\"\\nQuery: {query}\")\n",
    "        print(\"-\" * 50)\n",
    "\n",
    "        response = agent.create_turn(\n",
    "            messages=[{\"role\": \"user\", \"content\": query}],\n",
    "            session_id=session_id,\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": "markdown",
   "metadata": {},
   "source": [
    "Thanks for checking out this tutorial, hopefully you can now automate everything with Llama! :D\n",
    "\n",
    "Next up, we learn another hot topic of LLMs: Memory and Rag. Continue learning [here](./04_Memory101.ipynb)!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}