{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5715368",
   "metadata": {},
   "source": [
    "# Tracking token usage\n",
    "\n",
    "This notebook goes over how to track your token usage for specific calls. It is currently only implemented for the OpenAI API.\n",
    "\n",
    "Let's first look at an extremely simple example of tracking token usage for a single LLM call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9455db35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks import get_openai_callback\n",
    "from langchain.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1c55cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model_name=\"gpt-3.5-turbo-instruct\", n=2, best_of=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "31667d54",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 37\n",
      "\tPrompt Tokens: 4\n",
      "\tCompletion Tokens: 33\n",
      "Successful Requests: 1\n",
      "Total Cost (USD): $7.2e-05\n"
     ]
    }
   ],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    result = llm.invoke(\"Tell me a joke\")\n",
    "    print(cb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0ab6d27",
   "metadata": {},
   "source": [
    "Anything inside the context manager will get tracked. Here's an example of using it to track multiple calls in sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e09420f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "72\n"
     ]
    }
   ],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    result = llm.invoke(\"Tell me a joke\")\n",
    "    result2 = llm.invoke(\"Tell me a joke\")\n",
    "    print(cb.total_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8186e7b",
   "metadata": {},
   "source": [
    "If a chain or agent with multiple steps in it is used, it will track all those steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5d1125c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentType, initialize_agent, load_tools\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "llm = OpenAI(temperature=0)\n",
    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm)\n",
    "agent = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f98c536",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m I need to find out who Olivia Wilde's boyfriend is and then calculate his age raised to the 0.23 power.\n",
      "Action: Search\n",
      "Action Input: \"Olivia Wilde boyfriend\"\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m[\"Olivia Wilde and Harry Styles took fans by surprise with their whirlwind romance, which began when they met on the set of Don't Worry Darling.\", 'Olivia Wilde started dating Harry Styles after ending her years-long engagement to Jason Sudeikis — see their relationship timeline.', 'Olivia Wilde and Harry Styles were spotted early on in their relationship walking around London. (. Image ...', \"Looks like Olivia Wilde and Jason Sudeikis are starting 2023 on good terms. Amid their highly publicized custody battle – and the actress' ...\", 'The two started dating after Wilde split up with actor Jason Sudeikisin 2020. However, their relationship came to an end last November.', \"Olivia Wilde and Harry Styles started dating during the filming of Don't Worry Darling. While the movie got a lot of backlash because of the ...\", \"Here's what we know so far about Harry Styles and Olivia Wilde's relationship.\", 'Olivia and the Grammy winner kept their romance out of the spotlight as their relationship began just two months after her split from ex-fiancé ...', \"Harry Styles and Olivia Wilde first met on the set of Don't Worry Darling and stepped out as a couple in January 2021. Relive all their biggest relationship ...\"]\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m Harry Styles is Olivia Wilde's boyfriend.\n",
      "Action: Search\n",
      "Action Input: \"Harry Styles age\"\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m29 years\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I need to calculate 29 raised to the 0.23 power.\n",
      "Action: Calculator\n",
      "Action Input: 29^0.23\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3mAnswer: 2.169459462491557\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer.\n",
      "Final Answer: Harry Styles is Olivia Wilde's boyfriend and his current age raised to the 0.23 power is 2.169459462491557.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Total Tokens: 2205\n",
      "Prompt Tokens: 2053\n",
      "Completion Tokens: 152\n",
      "Total Cost (USD): $0.0441\n"
     ]
    }
   ],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    response = agent.run(\n",
    "        \"Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?\"\n",
    "    )\n",
    "    print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "    print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "    print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
    "    print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80ca77a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
