langchain/docs/docs/modules/model_io/chat/token_usage_tracking.ipynb

{
 "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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a55e87a-3291-4e7f-8e8e-4c69b0854384",
   "metadata": {},
   "source": [
    "## Using AIMessage.response_metadata\n",
    "\n",
    "A number of model providers return token usage information as part of the chat generation response. When available, this is included in the [AIMessage.response_metadata](/docs/modules/model_io/chat/response_metadata/). Here's an example with OpenAI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "467ccdeb-6b62-45e5-816e-167cd24d2586",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'token_usage': {'completion_tokens': 225,\n",
       "  'prompt_tokens': 17,\n",
       "  'total_tokens': 242},\n",
       " 'model_name': 'gpt-4-turbo',\n",
       " 'system_fingerprint': 'fp_76f018034d',\n",
       " 'finish_reason': 'stop',\n",
       " 'logprobs': None}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# !pip install -qU langchain-openai\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-turbo\")\n",
    "msg = llm.invoke([(\"human\", \"What's the oldest known example of cuneiform\")])\n",
    "msg.response_metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d5026e9-3ad4-41e6-9946-9f1a26f4a21f",
   "metadata": {},
   "source": [
    "And here's an example with Anthropic:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "145404f1-e088-4824-b468-236c486a9903",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'msg_01P61rdHbapEo6h3fjpfpCQT',\n",
       " 'model': 'claude-3-sonnet-20240229',\n",
       " 'stop_reason': 'end_turn',\n",
       " 'stop_sequence': None,\n",
       " 'usage': {'input_tokens': 17, 'output_tokens': 306}}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# !pip install -qU langchain-anthropic\n",
    "\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "\n",
    "llm = ChatAnthropic(model=\"claude-3-sonnet-20240229\")\n",
    "msg = llm.invoke([(\"human\", \"What's the oldest known example of cuneiform\")])\n",
    "msg.response_metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6845407-af25-4eed-bc3e-50925c6661e0",
   "metadata": {},
   "source": [
    "## Using callbacks\n",
    "\n",
    "There are also some API-specific callback context managers that allow you to track token usage across multiple calls. It is currently only implemented for the OpenAI API and Bedrock Anthropic API.\n",
    "\n",
    "### OpenAI\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": 3,
   "id": "9455db35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -qU langchain-community wikipedia\n",
    "\n",
    "from langchain_community.callbacks.manager import get_openai_callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d1c55cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-4-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31667d54",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 26\n",
      "\tPrompt Tokens: 11\n",
      "\tCompletion Tokens: 15\n",
      "Successful Requests: 1\n",
      "Total Cost (USD): $0.00056\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": 6,
   "id": "e09420f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "52\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": 17,
   "id": "5d1125c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent, load_tools\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You're a helpful assistant\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "tools = load_tools([\"wikipedia\"])\n",
    "agent = create_tool_calling_agent(llm, tools, prompt)\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent, tools=tools, verbose=True, stream_runnable=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1ae74d-8300-4041-9ff4-66093ee592b1",
   "metadata": {},
   "source": [
    "```{=mdx}\n",
    ":::note\n",
    "We have to set `stream_runnable=False` for token counting to work. By default the AgentExecutor will stream the underlying agent so that you can get the most granular results when streaming events via AgentExecutor.stream_events. However, OpenAI does not return token counts when streaming model responses, so we need to turn off the underlying streaming.\n",
    ":::\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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: `wikipedia` with `Hummingbird`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mPage: Hummingbird\n",
      "Summary: Hummingbirds are birds native to the Americas and comprise the biological family Trochilidae. With approximately 366 species and 113 genera, they occur from Alaska to Tierra del Fuego, but most species are found in Central and South America. As of 2024, 21 hummingbird species are listed as endangered or critically endangered, with numerous species declining in population.Hummingbirds have varied specialized characteristics to enable rapid, maneuverable flight: exceptional metabolic capacity, adaptations to high altitude, sensitive visual and communication abilities, and long-distance migration in some species. Among all birds, male hummingbirds have the widest diversity of plumage color, particularly in blues, greens, and purples. Hummingbirds are the smallest mature birds, measuring 7.5–13 cm (3–5 in) in length. The smallest is the 5 cm (2.0 in) bee hummingbird, which weighs less than 2.0 g (0.07 oz), and the largest is the 23 cm (9 in) giant hummingbird, weighing 18–24 grams (0.63–0.85 oz). Noted for long beaks, hummingbirds are specialized for feeding on flower nectar, but all species also consume small insects.\n",
      "They are known as hummingbirds because of the humming sound created by their beating wings, which flap at high frequencies audible to other birds and humans. They hover at rapid wing-flapping rates, which vary from around 12 beats per second in the largest species to 80 per second in small hummingbirds.\n",
      "Hummingbirds have the highest mass-specific metabolic rate of any homeothermic animal. To conserve energy when food is scarce and at night when not foraging, they can enter torpor, a state similar to hibernation, and slow their metabolic rate to 1⁄15 of its normal rate. While most hummingbirds do not migrate, the rufous hummingbird has one of the longest migrations among birds, traveling twice per year between Alaska and Mexico, a distance of about 3,900 miles (6,300 km).\n",
      "Hummingbirds split from their sister group, the swifts and treeswifts, around 42 million years ago. The oldest known fossil hummingbird is Eurotrochilus, from the Rupelian Stage of Early Oligocene Europe.\n",
      "\n",
      "\n",
      "\n",
      "Page: Bee hummingbird\n",
      "Summary: The bee hummingbird, zunzuncito or Helena hummingbird (Mellisuga helenae) is a species of hummingbird, native to the island of Cuba in the Caribbean. It is the smallest known bird. The bee hummingbird feeds on nectar of flowers and bugs found in Cuba.\n",
      "\n",
      "Page: Hummingbird cake\n",
      "Summary: Hummingbird cake is a banana-pineapple spice cake originating in Jamaica and a popular dessert in the southern United States since the 1970s. Ingredients include flour, sugar, salt, vegetable oil, ripe banana, pineapple, cinnamon, pecans, vanilla extract, eggs, and leavening agent. It is often served with cream cheese frosting.\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `wikipedia` with `Fastest bird`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mPage: Fastest animals\n",
      "Summary: This is a list of the fastest animals in the world, by types of animal.\n",
      "\n",
      "\n",
      "\n",
      "Page: List of birds by flight speed\n",
      "Summary: This is a list of the fastest flying birds in the world. A bird's velocity is necessarily variable; a hunting bird will reach much greater speeds while diving to catch prey than when flying horizontally. The bird that can achieve the greatest airspeed is the peregrine falcon, able to exceed 320 km/h (200 mph) in its dives. A close relative of the common swift, the white-throated needletail (Hirundapus caudacutus), is commonly reported as the fastest bird in level flight with a reported top speed of 169 km/h (105 mph). This record remains unconfirmed as the measurement methods have never been published or verified. The record for the fastest confirmed level flight by a bird is 111.5 km/h (69.3 mph) held by the common swift.\n",
      "\n",
      "Page: Ostrich\n",
      "Summary: Ostriches are large flightless birds. They are the heaviest and largest living birds, with adult common ostriches weighing anywhere between 63.5 and 145 kilograms and laying the largest eggs of any living land animal. With the ability to run at 70 km/h (43.5 mph), they are the fastest birds on land. They are farmed worldwide, with significant industries in the Philippines and in Namibia. Ostrich leather is a lucrative commodity, and the large feathers are used as plumes for the decoration of ceremonial headgear. Ostrich eggs have been used by humans for millennia.\n",
      "Ostriches are of the genus Struthio in the order Struthioniformes, part of the infra-class Palaeognathae, a diverse group of flightless birds also known as ratites that includes the emus, rheas, cassowaries, kiwis and the extinct elephant birds and moas. There are two living species of ostrich: the common ostrich, native to large areas of sub-Saharan Africa, and the Somali ostrich, native to the Horn of Africa.  The common ostrich was historically native to the Arabian Peninsula, and ostriches were present across Asia as far east as China and Mongolia during the Late Pleistocene and possibly into the Holocene.\u001b[0m\u001b[32;1m\u001b[1;3m### Hummingbird's Scientific Name\n",
      "The scientific name for the bee hummingbird, which is the smallest known bird and a species of hummingbird, is **Mellisuga helenae**. It is native to Cuba.\n",
      "\n",
      "### Fastest Bird Species\n",
      "The fastest bird in terms of airspeed is the **peregrine falcon**, which can exceed speeds of 320 km/h (200 mph) during its diving flight. In level flight, the fastest confirmed speed is held by the **common swift**, which can fly at 111.5 km/h (69.3 mph).\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Total Tokens: 1583\n",
      "Prompt Tokens: 1412\n",
      "Completion Tokens: 171\n",
      "Total Cost (USD): $0.019250000000000003\n"
     ]
    }
   ],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    response = agent_executor.invoke(\n",
    "        {\n",
    "            \"input\": \"What's a hummingbird's scientific name and what's the fastest bird species?\"\n",
    "        }\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": "markdown",
   "id": "ebc9122b-050b-4006-b763-264b0b26d9df",
   "metadata": {},
   "source": [
    "### Bedrock Anthropic\n",
    "\n",
    "The `get_bedrock_anthropic_callback` works very similarly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4a3eced5-2ff7-49a7-a48b-768af8658323",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 0\n",
      "\tPrompt Tokens: 0\n",
      "\tCompletion Tokens: 0\n",
      "Successful Requests: 2\n",
      "Total Cost (USD): $0.0\n"
     ]
    }
   ],
   "source": [
    "# !pip install langchain-aws\n",
    "from langchain_aws import ChatBedrock\n",
    "from langchain_community.callbacks.manager import get_bedrock_anthropic_callback\n",
    "\n",
    "llm = ChatBedrock(model_id=\"anthropic.claude-v2\")\n",
    "\n",
    "with get_bedrock_anthropic_callback() as cb:\n",
    "    result = llm.invoke(\"Tell me a joke\")\n",
    "    result2 = llm.invoke(\"Tell me a joke\")\n",
    "    print(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb40375d",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}