Harrison/openai callback (#684)

Author: hwchase17

docs/modules/agents/getting_started.ipynb

@@ -152,7 +152,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.9.0 64-bit ('llm-env')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },

docs/modules/llms/examples/token_usage_tracking.ipynb

@@ -0,0 +1,179 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e5715368",
+   "metadata": {},
+   "source": [
+    "# Token Usage Tracking\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.llms import OpenAI\n",
+    "from langchain.callbacks import get_openai_callback"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d1c55cc9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OpenAI(model_name=\"text-davinci-002\", n=2, best_of=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "31667d54",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "42\n"
+     ]
+    }
+   ],
+   "source": [
+    "with get_openai_callback() as cb:\n",
+    "    result = llm(\"Tell me a joke\")\n",
+    "    print(cb.total_tokens)"
+   ]
+  },
+  {
+   "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": [
+      "83\n"
+     ]
+    }
+   ],
+   "source": [
+    "with get_openai_callback() as cb:\n",
+    "    result = llm(\"Tell me a joke\")\n",
+    "    result2 = llm(\"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": 7,
+   "id": "5d1125c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.agents import load_tools\n",
+    "from langchain.agents import initialize_agent\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(tools, llm, agent=\"zero-shot-react-description\", verbose=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "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;3mJason Sudeikis\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m I need to find out Jason Sudeikis' age\n",
+      "Action: Search\n",
+      "Action Input: \"Jason Sudeikis age\"\u001b[0m\n",
+      "Observation: \u001b[36;1m\u001b[1;3m47 years\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m I need to calculate 47 raised to the 0.23 power\n",
+      "Action: Calculator\n",
+      "Action Input: 47^0.23\u001b[0m\n",
+      "Observation: \u001b[33;1m\u001b[1;3mAnswer: 2.4242784855673896\n",
+      "\u001b[0m\n",
+      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer\n",
+      "Final Answer: Jason Sudeikis, Olivia Wilde's boyfriend, is 47 years old and his age raised to the 0.23 power is 2.4242784855673896.\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n",
+      "1465\n"
+     ]
+    }
+   ],
+   "source": [
+    "with get_openai_callback() as cb:\n",
+    "    response = agent.run(\"Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?\")\n",
+    "    print(cb.total_tokens)"
+   ]
+  },
+  {
+   "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.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/modules/llms/generic_how_to.rst

@@ -9,6 +9,8 @@ The examples here all address certain "how-to" guides for working with LLMs.
 
 `Custom LLM <./examples/custom_llm.html>`_: How to create and use a custom LLM class, in case you have an LLM not from one of the standard providers (including one that you host yourself).
 
+`Token Usage Tracking <./examples/token_usage_tracking.html>`_: How to track the token usage of various chains/agents/LLM calls.
+
 
 .. toctree::
    :maxdepth: 1

langchain/callbacks/__init__.py

@@ -1,5 +1,9 @@
 """Callback handlers that allow listening to events in LangChain."""
+from contextlib import contextmanager
+from typing import Generator
+
 from langchain.callbacks.base import BaseCallbackHandler, BaseCallbackManager
+from langchain.callbacks.openai_info import OpenAICallbackHandler
 from langchain.callbacks.shared import SharedCallbackManager
 from langchain.callbacks.stdout import StdOutCallbackHandler
 
@@ -18,3 +22,13 @@ def set_handler(handler: BaseCallbackHandler) -> None:
 def set_default_callback_manager() -> None:
     """Set default callback manager."""
     set_handler(StdOutCallbackHandler())
+
+
+@contextmanager
+def get_openai_callback() -> Generator[OpenAICallbackHandler, None, None]:
+    """Get OpenAI callback handler in a context manager."""
+    handler = OpenAICallbackHandler()
+    manager = get_callback_manager()
+    manager.add_handler(handler)
+    yield handler
+    manager.remove_handler(handler)

langchain/callbacks/openai_info.py

@@ -0,0 +1,89 @@
+"""Callback Handler that prints to std out."""
+from typing import Any, Dict, List, Optional
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import AgentAction, AgentFinish, LLMResult
+
+
+class OpenAICallbackHandler(BaseCallbackHandler):
+    """Callback Handler that tracks OpenAI info."""
+
+    total_tokens: int = 0
+
+    @property
+    def always_verbose(self) -> bool:
+        """Whether to call verbose callbacks even if verbose is False."""
+        return True
+
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        """Print out the prompts."""
+        pass
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Do nothing."""
+        if response.llm_output is not None:
+            if "token_usage" in response.llm_output:
+                token_usage = response.llm_output["token_usage"]
+                if "total_tokens" in token_usage:
+                    self.total_tokens += token_usage["total_tokens"]
+
+    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Do nothing."""
+        pass
+
+    def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Print out that we are entering a chain."""
+        pass
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        """Print out that we finished a chain."""
+        pass
+
+    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+        """Do nothing."""
+        pass
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        action: AgentAction,
+        color: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Print out the log in specified color."""
+        pass
+
+    def on_tool_end(
+        self,
+        output: str,
+        color: Optional[str] = None,
+        observation_prefix: Optional[str] = None,
+        llm_prefix: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """If not the final action, print out observation."""
+        pass
+
+    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+        """Do nothing."""
+        pass
+
+    def on_text(
+        self,
+        text: str,
+        color: Optional[str] = None,
+        end: str = "",
+        **kwargs: Optional[str],
+    ) -> None:
+        """Run when agent ends."""
+        pass
+
+    def on_agent_finish(
+        self, finish: AgentFinish, color: Optional[str] = None, **kwargs: Any
+    ) -> None:
+        """Run on agent end."""
+        pass