{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import stackprinter # type: ignore\n", "import jupyter_black # type: ignore\n", "from dotenv import load_dotenv # type: ignore\n", "import time\n", "\n", "from baml_agents import init_logging, with_model\n", "from baml_client.async_client import b\n", "\n", "init_logging()\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": [ "# Using BAML in Jupyter notebokos\n", "\n", "## 1. Streaming outputs\n", "\n", "BAML offers a great developer experience, and I used some utilities to recreate some of the magic in Jupyter notebooks." ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from baml_agents.jupyter import JupyterOutputBox\n", "\n", "with JupyterOutputBox(clear_after_finish=False) as output_box:\n", " s = \"\"\n", " for c in \"abcdefghijklmnopqrstuvwxyz\" * 3:\n", " s += c\n", " output_box.update(s)\n", " time.sleep(0.008)\n", "\n", "# The letters are generated and displayed as they are being generated." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
{\n",
              "    \"genre\": \"Fantasy\",\n",
              "    \"characters\": [\n",
              "        {\n",
              "            \"name\": \"Lina\",\n",
              "            \"age\": 12,\n",
              "            \"occupation\": \"Apprentice Mage\"\n",
              "        },\n",
              "        {\n",
              "            \"name\": \"Eldon\",\n",
              "            \"age\": 45,\n",
              "            \"occupation\": \"Wizard\"\n",
              "        },\n",
              "        {\n",
              "            \"name\": \"Norik\",\n",
              "            \"age\": 17,\n",
              "            \"occupation\": \"Forest Ranger\"\n",
              "        }\n",
              "    ],\n",
              "    \"story_summary\": \"Lina seeks Eldon's help to unlock her magic potential; together, they face enchanted challenges in the Whispering Woods, discovering courage and friendship along the way.\"\n",
              "}
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from baml_agents.jupyter import JupyterBamlCollector\n", "\n", "with JupyterOutputBox(clear_after_finish=False) as s:\n", " c = JupyterBamlCollector(b, stream_callback=s.display(formatter=\"json\"))\n", " await c.b.InteractiveBamlJupyter_WriteShortStory()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "![Image](https://github.com/user-attachments/assets/ca21b62a-48b7-4a16-99fa-20bbe61f68b2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Inspecting prompts and completions" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", "
\n", " 

[system]
Write a short story about anything  Answer in JSON using this schema:
{
  genre: string,
  // at least three characters
  characters: [
    {
      name: string,
      age: int,
      occupation: string,
    }
  ],
  // Short story excerpt about anything, 30 words
  story_summary: string,
}

\n", "
\n", "\n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", "
\n", " 

{
  "genre": "Fantasy",
  "characters": [
    {
      "name": "Lina",
      "age": 12,
      "occupation": "Apprentice Mage"
    },
    {
      "name": "Eldon",
      "age": 45,
      "occupation": "Wizard"
    },
    {
      "name": "Norik",
      "age": 17,
      "occupation": "Forest Ranger"
    }
  ],
  "story_summary": "Lina seeks Eldon's help to unlock her magic potential; together, they face enchanted challenges in the Whispering Woods, discovering courage and friendship along the way."
}

\n", "
\n", "\n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "await c.display_calls()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "![Image](https://github.com/user-attachments/assets/59ed3cf3-84df-45a3-bf18-44dfbc457e4e)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can also group prompts and completions by LLM call and by LLM call session:" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", "
\n", " \n", "
\n", " \n", "
\n", " 

[system]
Write a short story about anything  Answer in JSON using this schema:
{
  genre: string,
  // at least three characters
  characters: [
    {
      name: string,
      age: int,
      occupation: string,
    }
  ],
  // Short story excerpt about anything, 30 words
  story_summary: string,
}

\n", "
\n", "
\n", " \n", " \n", "
\n", " \n", "
\n", " 

{
  "genre": "Fantasy",
  "characters": [
    {
      "name": "Liora",
      "age": 17,
      "occupation": "Young Mage"
    },
    {
      "name": "Eldric",
      "age": 45,
      "occupation": "Forest Guardian"
    },
    {
      "name": "Tom",
      "age": 12,
      "occupation": "Apprentice"
    }
  ],
  "story_summary": "Liora discovers a magical amulet in the enchanted forest, teaming up with Eldric and Tom to unlock its secrets and save their land from impending darkness."
}

\n", "
\n", "
\n", " \n", " \n", "
\n", " \n", "
\n", " 

[system]
Write a short story about anything  Answer in JSON using this schema:
{
  genre: string,
  // at least three characters
  characters: [
    {
      name: string,
      age: int,
      occupation: string,
    }
  ],
  // Short story excerpt about anything, 30 words
  story_summary: string,
}

\n", "
\n", "
\n", " \n", " \n", "
\n", " \n", "
\n", " 

{
  "genre": "Fantasy",
  "characters": [
    {
      "name": "Liora",
      "age": 17,
      "occupation": "Young Mage"
    },
    {
      "name": "Eldric",
      "age": 45,
      "occupation": "Forest Guardian"
    },
    {
      "name": "Tom",
      "age": 12,
      "occupation": "Apprentice"
    }
  ],
  "story_summary": "Liora discovers a magical amulet in the enchanted forest, teaming up with Eldric and Tom to unlock its secrets and save their land from impending darkness."
}

\n", "
\n", "
\n", " \n", " \n", "
\n", "\n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# Generate some calls\n", "c = JupyterBamlCollector(b)\n", "await c.b.InteractiveBamlJupyter_WriteShortStory()\n", "await c.b.InteractiveBamlJupyter_WriteShortStory()\n", "\n", "# Display the calls as a single session\n", "await c.display_session(\"Story generation session\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "![Image](https://github.com/user-attachments/assets/7bed8f67-efe9-48c7-918a-5d0853dded95)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Convenience class" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `JupyterOutputBox` and `JupyterBamlCollector` are available as a single convenience class:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
{'genre': 'Fantasy',\n",
              " 'characters': [{'name': 'Lina', 'age': 12, 'occupation': 'Apprentice Mage'}, {'name': 'Eldon', 'age': 45, 'occupation': 'Wizard'}, {'name': 'Norik', 'age': 17, 'occupation': 'Forest Ranger'}],\n",
              " 'story_summary': \"Lina seeks Eldon's help to unlock her magic potential; together, they face enchanted challenges in the Whispering Woods, discovering courage and friendship along the way.\"}
" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/javascript": "\n (function(){\n var el = document.getElementById(\"stream-8c1d5a3e-5935-4c61-a487-cd116b50a4de\");\n if (el) el.remove();\n })();\n ", "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", "
\n", " 

[system]
Write a short story about anything  Answer in JSON using this schema:
{
  genre: string,
  // at least three characters
  characters: [
    {
      name: string,
      age: int,
      occupation: string,
    }
  ],
  // Short story excerpt about anything, 30 words
  story_summary: string,
}

\n", "
\n", "\n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", "
\n", " 

{
  "genre": "Fantasy",
  "characters": [
    {
      "name": "Lina",
      "age": 12,
      "occupation": "Apprentice Mage"
    },
    {
      "name": "Eldon",
      "age": 45,
      "occupation": "Wizard"
    },
    {
      "name": "Norik",
      "age": 17,
      "occupation": "Forest Ranger"
    }
  ],
  "story_summary": "Lina seeks Eldon's help to unlock her magic potential; together, they face enchanted challenges in the Whispering Woods, discovering courage and friendship along the way."
}

\n", "
\n", "\n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from baml_agents.jupyter import JupyterBamlMonitor\n", "\n", "with JupyterBamlMonitor(b) as m:\n", " # Output is displayed in real time in the notebook\n", " await m.b.InteractiveBamlJupyter_WriteShortStory()\n", " # Outout is hidden from the notebook\n", "\n", "# Buttons with prompt and completino are displayed in the notebook\n", "await m.display_calls()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. Interactive chat" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Chat with your BAML Agent right in the notebook:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "a3e4af361e9740209105dfd34810f3bf", "version_major": 2, "version_minor": 0 }, "text/plain": [ "HBox(children=(Output(),), layout=Layout(display='inline-flex', flex_flow='column-reverse', max_height='500px'…" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e011be2694ce46fe93886aeaddb75c4e", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Box(children=(HTML(value='Loading...', layout=Layout(display='none')), Textarea(value='', continuous_update=Fa…" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from baml_agents.jupyter import ChatMessage, JupyterChatWidget\n", "\n", "\n", "async def callback(chat_history):\n", " msg = ChatMessage(\n", " content=f\"I'm seeing {len(chat_history)} messages\",\n", " role=\"assistant\",\n", " )\n", " return [msg]\n", "\n", "\n", "JupyterChatWidget(callback).display()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "![Image](https://github.com/user-attachments/assets/df46af62-97e5-42d6-8def-91762a46e432)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can also pre-fill chat:" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "0db671e549ef4b8da9a757c7fbe42b7c", "version_major": 2, "version_minor": 0 }, "text/plain": [ "HBox(children=(Output(outputs=({'output_type': 'display_data', 'data': {'text/plain': '\n", " document.querySelector(\".jupyter-button:last-child\").id = \"send-button\";\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "chat_history = [\n", " ChatMessage(\n", " content=\"Hey\",\n", " role=\"user\",\n", " ),\n", " ChatMessage(\n", " content=\"Hello\",\n", " role=\"assistant\",\n", " ),\n", "]\n", "\n", "JupyterChatWidget(callback, chat_history=chat_history).display()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can also invoke the callback automatically:" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "509d7a73188248ad8cf2c76c501a70e8", "version_major": 2, "version_minor": 0 }, "text/plain": [ "HBox(children=(Output(),), layout=Layout(display='inline-flex', flex_flow='column-reverse', max_height='500px'…" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "46d7d449927a415fb221f0cbdd8287b3", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Box(children=(HTML(value='Loading...', layout=Layout(display='none')), Textarea(value='', continuous_update=Fa…" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "chat = JupyterChatWidget(callback)\n", "await chat.run(\"It's a beautiful day\")\n", "# Callback is being run automatically on startup to respond to the \"It's a beautiful day\" message" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "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 }