From 3a30e6daa89735717a507c2fe4e9897358803f06 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Sun, 22 Jan 2023 23:37:01 -0800 Subject: [PATCH] Harrison/openai callback (#684) --- docs/modules/agents/getting_started.ipynb | 2 +- .../llms/examples/token_usage_tracking.ipynb | 179 ++++++++++++++++++ docs/modules/llms/generic_how_to.rst | 2 + langchain/callbacks/__init__.py | 14 ++ langchain/callbacks/openai_info.py | 89 +++++++++ 5 files changed, 285 insertions(+), 1 deletion(-) create mode 100644 docs/modules/llms/examples/token_usage_tracking.ipynb create mode 100644 langchain/callbacks/openai_info.py diff --git a/docs/modules/agents/getting_started.ipynb b/docs/modules/agents/getting_started.ipynb index 766d144a..1279ad83 100644 --- a/docs/modules/agents/getting_started.ipynb +++ b/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" }, diff --git a/docs/modules/llms/examples/token_usage_tracking.ipynb b/docs/modules/llms/examples/token_usage_tracking.ipynb new file mode 100644 index 00000000..199fcf7c --- /dev/null +++ b/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 +} diff --git a/docs/modules/llms/generic_how_to.rst b/docs/modules/llms/generic_how_to.rst index 7606eb33..90d1a837 100644 --- a/docs/modules/llms/generic_how_to.rst +++ b/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 diff --git a/langchain/callbacks/__init__.py b/langchain/callbacks/__init__.py index 13a4c302..501de6d3 100644 --- a/langchain/callbacks/__init__.py +++ b/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) diff --git a/langchain/callbacks/openai_info.py b/langchain/callbacks/openai_info.py new file mode 100644 index 00000000..b1471412 --- /dev/null +++ b/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