{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from dotenv import load_dotenv  # type: ignore\n",
    "from notebooks._utils import (\n",
    "    celsius_to_fahrenheit,\n",
    "    city_to_number,\n",
    "    city_to_weather_condition,\n",
    ")\n",
    "from pydantic import Field  # type: ignore\n",
    "from baml_agents import with_model\n",
    "from baml_client.sync_client import b\n",
    "import jupyter_black  # type: ignore\n",
    "import stackprinter  # type: ignore\n",
    "\n",
    "from baml_client import types\n",
    "from baml_client import types as T\n",
    "from baml_client.type_builder import TypeBuilder\n",
    "\n",
    "stackprinter.set_excepthook()\n",
    "load_dotenv()\n",
    "jupyter_black.load()\n",
    "\n",
    "b = with_model(b, \"gpt-4.1-nano\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Personally, I prefer the term \"Actions,\" so I'll use \"Tools\" and \"Actions\" interchangeably.\n",
    "\n",
    "Let's create a tool for an AI to use (Action). We'll be using Pydantic-model-as-a-function pattern rather than regular typed Python functions because that makes it easer to add parameter validation, type hints, and baml-specific (LLM-specific) information to the function parameters.\n",
    "\n",
    "So instead of writing...:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from baml_agents import Result\n",
    "\n",
    "\n",
    "def get_weather_info(\n",
    "    city: str,\n",
    "    measurement: Literal[\"celsius\", \"fahrenheit\"] = \"celsius\",\n",
    ") -> Result:\n",
    "    \"\"\"Get weather information for a given city.\"\"\"\n",
    "    return Result(\n",
    "        content=\"The weather in London is 18 degrees celsius with windy conditions.\",\n",
    "        error=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...we will do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from baml_agents import Action\n",
    "\n",
    "\n",
    "class GetWeatherInfo(Action):\n",
    "    \"\"\"Get weather information for a given city.\"\"\"\n",
    "\n",
    "    city: str\n",
    "    measurement: Literal[\"celsius\", \"fahrenheit\"] | None = None\n",
    "\n",
    "    def run(self) -> Result:\n",
    "        self.measurement = self.measurement or \"celsius\"\n",
    "        c = city_to_number(self.city, -10, 35)\n",
    "        condition = city_to_weather_condition(self.city)\n",
    "        if self.measurement.lower() == \"fahrenheit\":\n",
    "            c, u = celsius_to_fahrenheit(c), \"fahrenheit\"\n",
    "        else:\n",
    "            u = \"celsius\"\n",
    "        content = f\"The weather in {self.city} is {round(c, 1)} degrees {u} with {condition.lower()} conditions.\"\n",
    "        return Result(content=content, error=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try some cities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "content='The weather in New York is 57.2 degrees fahrenheit with foggy conditions.' error=False\n",
      "content='The weather in London is -7 degrees celsius with sunny conditions.' error=False\n",
      "content='The weather in Paris is -10 degrees celsius with rainy conditions.' error=False\n"
     ]
    }
   ],
   "source": [
    "print(GetWeatherInfo(city=\"New York\", measurement=\"fahrenheit\").run())\n",
    "print(GetWeatherInfo(city=\"London\").run())\n",
    "print(GetWeatherInfo(city=\"Paris\").run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manually configuring baml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "According to [12 Factor Agents](https://github.com/humanlayer/12-factor-agents), [tools are simply structured outputs](https://github.com/search?q=repo%3Ahumanlayer%2F12-factor-agents+Tools+are+just+structured+outputs&type=code).\n",
    "\n",
    "That's why we'll just create a baml prompt template:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class BamlCustomTools_GetWeatherInfo_Parameters {\n",
      "    city string\n",
      "    measurement \"celsius\" | \"fahrenheit\" @description(#\"default: celsius\"#)\n",
      "}\n",
      "function BamlCustomTools_GetWeatherInfo(text: string) -> BamlCustomTools_GetWeatherInfo_Parameters {\n",
      "  client Default\n",
      "  prompt #\"\n",
      "    {{ text }}\n",
      "    {{ ctx.output_format }}\n",
      "  \"#\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# This command shows file contents\n",
    "!awk '/^class BamlCustomTools_GetWeatherInfo_Parameters/,/^}/' ../baml_src/notebooks/02_baml_custom_tools.baml\n",
    "!awk '/^function BamlCustomTools_GetWeatherInfo/,/^}/' ../baml_src/notebooks/02_baml_custom_tools.baml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see the actual prompt that we will send to the LLM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "What's the weather like in the capital of Japan?\n",
      "Answer in JSON using this schema:\n",
      "{\n",
      "  city: string,\n",
      "  // default: celsius\n",
      "  measurement: \"celsius\" or \"fahrenheit\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import display_prompt\n",
    "\n",
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "\n",
    "request = b.request.BamlCustomTools_GetWeatherInfo(question)\n",
    "display_prompt(request)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's send it and run the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "city='Tokyo' measurement='celsius'\n"
     ]
    }
   ],
   "source": [
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "\n",
    "# Calling LLM\n",
    "completion = b.BamlCustomTools_GetWeatherInfo(question)\n",
    "print(completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Result(content='The weather in Tokyo is 26 degrees celsius with rainy conditions.', error=False)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Running the action\n",
    "GetWeatherInfo.validate(completion).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatically configuring baml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be using the [dynamic features of baml](https://docs.boundaryml.com/guide/baml-advanced/dynamic-runtime-types)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class BamlCustomTools_NextAction {\n",
      "    @@dynamic\n",
      "}\n",
      "function BamlCustomTools_GetNextAction(goal: string) -> BamlCustomTools_NextAction {\n",
      "    client Default\n",
      "    prompt #\"\n",
      "        Please select the next best action to achieve the goal.\n",
      "        \n",
      "        <goal>\n",
      "        {{ goal }}\n",
      "        </goal>\n",
      "        \n",
      "        {{ ctx.output_format(or_splitter=\" OR \", prefix=\"Answer in JSON format with exactly one next best action:\\n\") }}\n",
      "    \"#\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# This command shows file contents\n",
    "!awk '/^class BamlCustomTools_NextAction/,/^}/' ../baml_src/notebooks/02_baml_custom_tools.baml\n",
    "!awk '/^function BamlCustomTools_GetNextAction/,/^}/' ../baml_src/notebooks/02_baml_custom_tools.baml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To add fields from the Tool to a class, we must first know **which** class we want to add the fields to. This can be done in two ways:\n",
    "#### Choosing Target Class - Method 1: Simple\n",
    "\n",
    "It works by taking a look at the baml function annotation return type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "Please select the next best action to achieve the goal.\n",
      "\n",
      "‹goal›\n",
      "What's the weather like in the capital of Japan?\n",
      "‹/goal›\n",
      "\n",
      "Answer in JSON format with exactly one next best action:\n",
      "{\n",
      "  chosen_action: {\n",
      "    // Get weather information for a given city.\n",
      "    action_id: \"get_weather_info\",\n",
      "    city: string,\n",
      "    measurement: \"celsius\" OR \"fahrenheit\" OR null,\n",
      "  },\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import ActionRunner\n",
    "\n",
    "r = ActionRunner(TypeBuilder, b=b)\n",
    "r.add_action(GetWeatherInfo)\n",
    "\n",
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "request = r.b.request.BamlCustomTools_GetNextAction(question)\n",
    "\n",
    "display_prompt(request)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Choosing Target Class - Method 2\n",
    "\n",
    "You will only need it if you're adding tools not to the root class, otherwise, skip this\n",
    "\n",
    "... [code moved to the bottom of the file]\n",
    "\n",
    "#### Choosing Target Class - Method 3\n",
    "\n",
    "You will only need this so you could fall back to a simpler method if you encouter bugs or need more customization\n",
    "\n",
    "... [code moved to the bottom of the file]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Let's send the request, get the completion and run the action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chosen_action={'action_id': 'get_weather_info', 'city': 'Tokyo', 'measurement': 'celsius'}\n"
     ]
    }
   ],
   "source": [
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "\n",
    "r = ActionRunner(TypeBuilder)\n",
    "r.add_action(GetWeatherInfo)\n",
    "\n",
    "# Calling the LLM\n",
    "action = b.BamlCustomTools_GetNextAction(\n",
    "    question, baml_options={\"tb\": r.tb(T.BamlCustomTools_NextAction)}\n",
    ")\n",
    "print(action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Result(content='The weather in Tokyo is 26 degrees celsius with rainy conditions.', error=False)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Running the action chosen by LLM\n",
    "r.run(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix A: Advanced Customization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can safely skip this section, you won't need this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Choosing Target Class - Method 2\n",
    "\n",
    "You will only need it if you're adding tools not to the root class, otherwise, skip this\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "Please select the next best action to achieve the goal.\n",
      "\n",
      "‹goal›\n",
      "What's the weather like in the capital of Japan?\n",
      "‹/goal›\n",
      "\n",
      "Answer in JSON format with exactly one next best action:\n",
      "{\n",
      "  chosen_action: {\n",
      "    // Get weather information for a given city.\n",
      "    action_id: \"get_weather_info\",\n",
      "    city: string,\n",
      "    measurement: \"celsius\" OR \"fahrenheit\" OR null,\n",
      "  },\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import ActionRunner\n",
    "\n",
    "r = ActionRunner(TypeBuilder, b=b)\n",
    "r.add_action(GetWeatherInfo)\n",
    "\n",
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "request = r.b_(return_class=\"BamlCustomTools_NextAction\").BamlCustomTools_GetNextAction(\n",
    "    question\n",
    ")\n",
    "# OR\n",
    "rc = types.BamlCustomTools_NextAction\n",
    "request = r.b_(return_class=rc).request.BamlCustomTools_GetNextAction(question)\n",
    "\n",
    "display_prompt(request)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Choosing Target Class - Method 3\n",
    "\n",
    "You will only need this so you could fall back to a simpler method if you encouter bugs or need more customization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "Please select the next best action to achieve the goal.\n",
      "\n",
      "‹goal›\n",
      "What's the weather like in the capital of Japan?\n",
      "‹/goal›\n",
      "\n",
      "Answer in JSON format with exactly one next best action:\n",
      "{\n",
      "  chosen_action: {\n",
      "    // Get weather information for a given city.\n",
      "    action_id: \"get_weather_info\",\n",
      "    city: string,\n",
      "    measurement: \"celsius\" OR \"fahrenheit\" OR null,\n",
      "  },\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import ActionRunner\n",
    "\n",
    "r = ActionRunner(TypeBuilder)\n",
    "r.add_action(GetWeatherInfo)\n",
    "\n",
    "# Choose one of the following ways to create the type builder:\n",
    "tb = r.tb(\"BamlCustomTools_NextAction\")  # Either way works\n",
    "tb = r.tb(\"BamlCustomTools_NextAction\", tb=TypeBuilder())  # Either way works\n",
    "tb = r.tb(types.BamlCustomTools_NextAction)  # Either way works\n",
    "tb = r.tb(types.BamlCustomTools_NextAction, tb=TypeBuilder())  # Either way works\n",
    "\n",
    "question = \"What's the weather like in the capital of Japan?\"\n",
    "request = b.request.BamlCustomTools_GetNextAction(question, baml_options={\"tb\": tb})\n",
    "display_prompt(request)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}