{
 "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 Chat model call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9455db35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks import get_openai_callback\n",
    "from langchain.chat_models import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d1c55cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model_name=\"gpt-4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "31667d54",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 24\n",
      "\tPrompt Tokens: 11\n",
      "\tCompletion Tokens: 13\n",
      "Successful Requests: 1\n",
      "Total Cost (USD): $0.0011099999999999999\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": 10,
   "id": "e09420f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48\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": 11,
   "id": "5d1125c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentType, initialize_agent, load_tools\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm)\n",
    "agent = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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\n",
      "Invoking: `Search` with `Olivia Wilde's current boyfriend`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m['Things are looking golden for Olivia Wilde, as the actress has jumped back into the dating pool following her split from Harry Styles — read ...', \"“I did not want service to take place at the home of Olivia's current partner because Otis and Daisy might be present,” Sudeikis wrote in his ...\", \"February 2021: Olivia Wilde praises Harry Styles' modesty. One month after the duo made headlines with their budding romance, Wilde gave her new beau major ...\", 'An insider revealed to People that the new couple had been dating for some time. \"They were in Montecito, California this weekend for a wedding, ...', 'A source told People last year that Wilde and Styles were still friends despite deciding to take a break. \"He\\'s still touring and is now going ...', \"... love life. “He's your typical average Joe.” The source adds, “She's not giving too much away right now and wants to keep the relationship ...\", \"Multiple sources said the two were “taking a break” from dating because of distance and different priorities. “He's still touring and is now ...\", 'Comments. Filed under. celebrity couples · celebrity dating · harry styles · jason sudeikis · olivia wilde ... Now Holds A Darker MeaningNYPost.', '... dating during filming. The 39-year-old did however look very cosy with the comedian, although his relationship status is unknown. Olivia ...']\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Search` with `Harry Styles current age`\n",
      "responded: Olivia Wilde's current boyfriend is Harry Styles. Let me find out his age for you.\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m29 years\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Calculator` with `29 ^ 0.23`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mAnswer: 2.169459462491557\u001b[0m\u001b[32;1m\u001b[1;3mHarry Styles' current age (29 years) raised to the 0.23 power is approximately 2.17.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Total Tokens: 1929\n",
      "Prompt Tokens: 1799\n",
      "Completion Tokens: 130\n",
      "Total Cost (USD): $0.06176999999999999\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}\")"
   ]
  }
 ],
 "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
}
