| { |
| "cells": [ |
| { |
| "cell_type": "code", |
| "id": "8ps8bnh6lo3", |
| "source": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements. See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership. The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License. You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied. See the License for the\n# specific language governing permissions and limitations\n# under the License.", |
| "metadata": {}, |
| "execution_count": null, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "5eebfc21-b54a-435d-84dc-b9de1d620e4e", |
| "metadata": {}, |
| "source": [ |
| "# Tool Calling\n", |
| "\n", |
| "In this example we'll go over a basic application that selects between a few different tools, calls them, and uses an LLM to process the results.\n", |
| "\n", |
| "Requirements:\n", |
| "- OpenAI API key, set as the env variable `OPENAI_API_KEY`\n", |
| "- [tomorrow.io](tomorrow.io) API Key, set as the env variable `TOMORROW_API_KEY`" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "537a0463-1d35-48dd-b0af-9ebabf2b08dc", |
| "metadata": {}, |
| "source": [ |
| "# Imports" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 5, |
| "id": "b60de9a5-43b9-421d-b2c4-a758d3fcde3f", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "import inspect\n", |
| "import json\n", |
| "import os\n", |
| "from typing import Callable, Optional\n", |
| "\n", |
| "import openai\n", |
| "import requests\n", |
| "\n", |
| "from burr.core import State, action, when\n", |
| "from burr.core.application import ApplicationBuilder\n", |
| "\n", |
| "# Set up + instantiate instrumentation\n", |
| "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n", |
| "\n", |
| "OpenAIInstrumentor().instrument()" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "5cc27a92-e36f-4306-91b8-aa68f7d2e6a2", |
| "metadata": {}, |
| "source": [ |
| "# Defining tools\n", |
| "\n", |
| "Let's define a few tools that we want to use. They'll all have:\n", |
| "\n", |
| "1. Primitive input, annotated with types\n", |
| "2. return types of dictionaries (free-form)\n", |
| "\n", |
| "This way we can have some structure in the return type. You'll likely want a stricter return type for your cases" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 2, |
| "id": "690935f7-19b4-494d-abab-9a3f7cf0c26c", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "def _weather_tool(latitude: float, longitude: float) -> dict:\n", |
| " \"\"\"Queries the weather for a given latitude and longitude.\"\"\"\n", |
| " api_key = os.environ.get(\"TOMORROW_API_KEY\")\n", |
| " url = f\"https://api.tomorrow.io/v4/weather/forecast?location={latitude},{longitude}&apikey={api_key}\"\n", |
| "\n", |
| " response = requests.get(url)\n", |
| " if response.status_code == 200:\n", |
| " return response.json()\n", |
| " else:\n", |
| " return {\"error\": f\"Failed to get weather data. Status code: {response.status_code}\"}\n", |
| "\n", |
| "\n", |
| "def _text_wife_tool(message: str) -> dict:\n", |
| " \"\"\"Texts your wife with a message.\"\"\"\n", |
| " # Dummy implementation for the text wife tool\n", |
| " # Replace this with actual SMS API logic\n", |
| " return {\"action\": f\"Texted wife: {message}\"}\n", |
| "\n", |
| "\n", |
| "def _order_coffee_tool(\n", |
| " size: str, coffee_preparation: str, any_modifications: Optional[str] = None\n", |
| ") -> dict:\n", |
| " \"\"\"Orders a coffee with the given size, preparation, and any modifications.\"\"\"\n", |
| " # Dummy implementation for the order coffee tool\n", |
| " # Replace this with actual coffee shop API logic\n", |
| " return {\n", |
| " \"action\": (\n", |
| " f\"Ordered a {size} {coffee_preparation}\" + \"with {any_modifications}\"\n", |
| " if any_modifications\n", |
| " else \"\"\n", |
| " )\n", |
| " }\n", |
| "\n", |
| "\n", |
| "def _fallback(response: str) -> dict:\n", |
| " \"\"\"Tells the user that the assistant can't do that -- this should be a fallback\"\"\"\n", |
| " return {\"response\": response}\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "5c198e8d-a467-433c-92f0-b125d3529750", |
| "metadata": {}, |
| "source": [ |
| "# Defining the tools (as OpenAI wants them)\n", |
| "\n", |
| "We do a little cleverness to get types -- you'll likely want to use pydantic to get this right later on:" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 8, |
| "id": "884bf319-5648-4049-a654-ef2305f32ec0", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "# You'll want to add more types here as needed\n", |
| "TYPE_MAP = {\n", |
| " str: \"string\",\n", |
| " int: \"integer\",\n", |
| " float: \"number\",\n", |
| " bool: \"boolean\",\n", |
| "}\n", |
| "\n", |
| "# You can also consider using a library like pydantic to further integrate with OpenAI\n", |
| "OPENAI_TOOLS = [\n", |
| " {\n", |
| " \"type\": \"function\",\n", |
| " \"function\": {\n", |
| " \"name\": fn_name,\n", |
| " \"description\": fn.__doc__ or fn_name,\n", |
| " \"parameters\": {\n", |
| " \"type\": \"object\",\n", |
| " \"properties\": {\n", |
| " param.name: {\n", |
| " \"type\": TYPE_MAP.get(param.annotation)\n", |
| " or \"string\", # TODO -- add error cases\n", |
| " \"description\": param.name,\n", |
| " }\n", |
| " for param in inspect.signature(fn).parameters.values()\n", |
| " },\n", |
| " \"required\": [param.name for param in inspect.signature(fn).parameters.values()],\n", |
| " },\n", |
| " },\n", |
| " }\n", |
| " for fn_name, fn in {\n", |
| " \"query_weather\": _weather_tool,\n", |
| " \"order_coffee\": _order_coffee_tool,\n", |
| " \"text_wife\": _text_wife_tool,\n", |
| " \"fallback\": _fallback,\n", |
| " }.items()\n", |
| "]" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "90c27937-e723-4f4f-ad5f-70bc518919ba", |
| "metadata": {}, |
| "source": [ |
| "# Defining our actions\n", |
| "\n", |
| "We're going to have:\n", |
| "1. An action that processes the user input\n", |
| "2. An action that selects the tool + parameters\n", |
| "3. An action that calls the tool\n", |
| "4. An action that formats the results\n", |
| "\n", |
| "This way each is decoupled. The action that calls the tool will be \"generic\", and we'll be binding it to each tool (see next section)" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 3, |
| "id": "c45bea6d-f5af-4144-8d77-3d631d127b6e", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "@action(reads=[], writes=[\"query\"])\n", |
| "def process_input(state: State, query: str) -> State:\n", |
| " \"\"\"Simple action to process input for the assistant.\"\"\"\n", |
| " return state.update(query=query)\n", |
| "\n", |
| "\n", |
| "\n", |
| "@action(reads=[\"query\"], writes=[\"tool_parameters\", \"tool\"])\n", |
| "def select_tool(state: State) -> State:\n", |
| " \"\"\"Selects the tool + assigns the parameters. Uses the tool-calling API.\"\"\"\n", |
| " messages = [\n", |
| " {\n", |
| " \"role\": \"system\",\n", |
| " \"content\": (\n", |
| " \"You are a helpful assistant. Use the supplied tools to assist the user, if they apply in any way. Remember to use the tools! They can do stuff you can't.\"\n", |
| " \"If you can't use only the tools provided to answer the question but know the answer, please provide the answer\"\n", |
| " \"If you cannot use the tools provided to answer the question, use the fallback tool and provide a reason. \"\n", |
| " \"Again, if you can't use one tool provided to answer the question, use the fallback tool and provide a reason. \"\n", |
| " \"You must select exactly one tool no matter what, filling in every parameters with your best guess. Do not skip out on parameters!\"\n", |
| " ),\n", |
| " },\n", |
| " {\"role\": \"user\", \"content\": state[\"query\"]},\n", |
| " ]\n", |
| " response = openai.chat.completions.create(\n", |
| " model=\"gpt-4o\",\n", |
| " messages=messages,\n", |
| " tools=OPENAI_TOOLS,\n", |
| " )\n", |
| "\n", |
| " # Extract the tool name and parameters from OpenAI's response\n", |
| " if response.choices[0].message.tool_calls is None:\n", |
| " return state.update(\n", |
| " tool=\"fallback\",\n", |
| " tool_parameters={\n", |
| " \"response\": \"No tool was selected, instead response was: {response.choices[0].message}.\"\n", |
| " },\n", |
| " )\n", |
| " fn = response.choices[0].message.tool_calls[0].function\n", |
| "\n", |
| " return state.update(tool=fn.name, tool_parameters=json.loads(fn.arguments))\n", |
| "\n", |
| "\n", |
| "@action(reads=[\"tool_parameters\"], writes=[\"raw_response\"])\n", |
| "def call_tool(state: State, tool_function: Callable) -> State:\n", |
| " \"\"\"Action to call the tool. This will be bound to the tool function.\"\"\"\n", |
| " response = tool_function(**state[\"tool_parameters\"])\n", |
| " return state.update(raw_response=response)\n", |
| "\n", |
| "\n", |
| "@action(reads=[\"query\", \"raw_response\"], writes=[\"final_output\"])\n", |
| "def format_results(state: State) -> State:\n", |
| " \"\"\"Action to format the results in a usable way. Note we're not cascading in context for the chat history.\n", |
| " This is largely due to keeping it simple, but you'll likely want to pass IDs around or maintain the chat history yourself\n", |
| " \"\"\"\n", |
| " response = openai.chat.completions.create(\n", |
| " model=\"gpt-4o\",\n", |
| " messages=[\n", |
| " {\n", |
| " \"role\": \"system\",\n", |
| " \"content\": (\n", |
| " \"You are a helpful assistant. Your goal is to take the\"\n", |
| " \"data presented and use it to answer the original question:\"\n", |
| " ),\n", |
| " },\n", |
| " {\n", |
| " \"role\": \"user\",\n", |
| " \"content\": (\n", |
| " f\"The original question was: {state['query']}.\"\n", |
| " f\"The data is: {state['raw_response']}. Please format\"\n", |
| " \"the data and provide a response that responds to the original query.\"\n", |
| " \"As always, be concise (tokens aren't free!).\"\n", |
| " ),\n", |
| " },\n", |
| " ],\n", |
| " )\n", |
| "\n", |
| " return state.update(final_output=response.choices[0].message.content)" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "fccd23db-b8d6-4179-a452-b37defbeff51", |
| "metadata": {}, |
| "source": [ |
| "# Building the application\n", |
| "\n", |
| "Now let's build the application. We will:\n", |
| "1. Add the right actions\n", |
| "2. Add the right transitions\n", |
| "3. Set up a tracker" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 4, |
| "id": "9477c4f1-96ad-4125-a504-ab47ac4e56d2", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "app = (\n", |
| " ApplicationBuilder()\n", |
| " .with_actions(\n", |
| " process_input,\n", |
| " select_tool,\n", |
| " format_results,\n", |
| " query_weather=call_tool.bind(tool_function=_weather_tool),\n", |
| " text_wife=call_tool.bind(tool_function=_text_wife_tool),\n", |
| " order_coffee=call_tool.bind(tool_function=_order_coffee_tool),\n", |
| " fallback=call_tool.bind(tool_function=_fallback),\n", |
| " )\n", |
| " .with_transitions(\n", |
| " (\"process_input\", \"select_tool\"),\n", |
| " (\"select_tool\", \"query_weather\", when(tool=\"query_weather\")),\n", |
| " (\"select_tool\", \"text_wife\", when(tool=\"text_wife\")),\n", |
| " (\"select_tool\", \"order_coffee\", when(tool=\"order_coffee\")),\n", |
| " (\"select_tool\", \"fallback\", when(tool=\"fallback\")),\n", |
| " ([\"query_weather\", \"text_wife\", \"order_coffee\", \"fallback\"], \"format_results\"),\n", |
| " (\"format_results\", \"process_input\"),\n", |
| " )\n", |
| " .with_entrypoint(\"process_input\")\n", |
| " .with_tracker(project=\"test_tool_calling\", use_otel_tracing=True)\n", |
| " .build()\n", |
| ")" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "8b7d206e-0c4f-40d2-a4f1-21341f8d26b6", |
| "metadata": {}, |
| "source": [ |
| "# Running the app\n", |
| "\n", |
| "And it works! " |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 9, |
| "id": "8b94805e-6da2-4fbe-aa05-92acd1e63ca4", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "Based on the data provided for San Francisco as of September 30, 2024, and subsequent days:\n", |
| "\n", |
| "- **Temperature**: Currently around 25.5°C to 25.9°C.\n", |
| "- **Weather**: Clear skies, with cloud cover at 0%.\n", |
| "- **Humidity**: Ranges between 35% to 46%.\n", |
| "- **Wind**: Light winds from the north, shifting to the west, with speeds ranging from 1.63 to 3.48 km/h. Gusts can reach up to 4.17 km/h.\n", |
| "- **Visibility**: Excellent, at 16 km.\n", |
| "- **Pressure**: Around 1007.7 hPa.\n", |
| "- **UV Index**: Moderate to high, peaking at 6.\n", |
| "\n", |
| "This indicates a clear, warm day with comfortable humidity levels and light winds.\n" |
| ] |
| } |
| ], |
| "source": [ |
| "app.visualize(output_file_path=\"./statemachine.png\")\n", |
| "action, result, state = app.run(\n", |
| " halt_after=[\"format_results\"],\n", |
| " inputs={\"query\": \"What's the weather like in San Francisco?\"},\n", |
| ")\n", |
| "print(state[\"final_output\"])" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "8dd763c1-128b-4b78-a973-ac01723ec636", |
| "metadata": {}, |
| "source": [ |
| "# Visualizing\n", |
| "\n", |
| "We can view in the UI. Make sure `burr` is runnng for this link to work" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": null, |
| "id": "6764124d-eb0f-43f7-92e7-808893b9a831", |
| "metadata": {}, |
| "outputs": [], |
| "source": [ |
| "from IPython.display import Markdown\n", |
| "url = f\"[Link to UI](http://localhost:7241/project/demo_email_assistant/{app.uid})\"\n", |
| "Markdown(url)" |
| ] |
| } |
| ], |
| "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.12.0" |
| } |
| }, |
| "nbformat": 4, |
| "nbformat_minor": 5 |
| } |
