LangChainでRAGやエージェントを組んでいると、思ったより使用トークン量がおおくなってしまうことってありますよね。
でも、いまいち使用している総量がわからない、いちいちopenai.comのusage見るのも面倒だな、という状況だったので、invokeのCallbacks機能を使用して、トータルでどれくらい使用されたのかを追跡してみました。
ポイントは、簡単に、もとのソースコードを極力汚さない、です。
コストの計算にはトークン追跡が必須なので、是非参照ください。
動作環境
- Windows 11
- WSL2
- Ubuntu 22.04
$ uv --version
uv 0.5.1
$ uv run python -V
Python 3.12.7
$ uv pip freeze
annotated-types==0.7.0
anyio==4.8.0
certifi==2025.1.31
charset-normalizer==3.4.1
distro==1.9.0
h11==0.14.0
httpcore==1.0.7
httpx==0.28.1
idna==3.10
jiter==0.8.2
jsonpatch==1.33
jsonpointer==3.0.0
langchain-core==0.3.33
langchain-openai==0.3.3
langgraph==0.2.69
langgraph-checkpoint==2.0.10
langgraph-sdk==0.1.51
langsmith==0.3.4
msgpack==1.1.0
openai==1.61.0
orjson==3.10.15
packaging==24.2
pydantic==2.10.6
pydantic-core==2.27.2
python-dotenv==1.0.1
pyyaml==6.0.2
regex==2024.11.6
requests==2.32.3
requests-toolbelt==1.0.0
sniffio==1.3.1
tenacity==9.0.0
tiktoken==0.8.0
tqdm==4.67.1
typing-extensions==4.12.2
urllib3==2.3.0
zstandard==0.23.0※ 他の検証でも使用しているので、不要なライブラリが多々はいっています。
コスト追跡関数の実装方法
今回はLangChainのCallbacksは、LLMの様々な実行状態時に呼ばれる関数です。
詳細はほかサイトや、下記公式サイトを参照ください。
Callbacks | 🦜️🔗 LangChain
Head to Integrations for documentation on built-in callbacks integrations with 3rd-party tools.
python.langchain.com
今回は、BaseCallbackHandlerクラスの中の、下記をオーバーライドして使います。
こちらは、コメントにある通り、LLMの実行終了時に呼ばれる関数です。
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
"""Run when LLM ends running."""さて、こちらがコード全文となります。
簡単な質問を1回、生成AIに聞いてみるコードです。
なお、OPENAI_API_KEYは.envファイルに記載したのを読み込んでいます
import os
from typing import Any
from dotenv import load_dotenv
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_core.output_parsers import StrOutputParser
from langchain_core.outputs.llm_result import LLMResult
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import ConfigurableField
from langchain_openai import ChatOpenAI
load_dotenv()
# BaseCallbackHandlerを継承したカスタムコールバックハンドラ
class CountTokensHandler(BaseCallbackHandler):
def __init__(self):
self.input_tokens = 0
self.output_tokens = 0
self.total_tokens = 0
self.invoke_count = 0
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
self.invoke_count += 1
token_usage = (
response.llm_output.get("token_usage", {})
if response.llm_output is not None
else {}
)
self.input_tokens += token_usage.get("prompt_tokens", 0)
self.output_tokens += token_usage.get("completion_tokens", 0)
self.total_tokens += token_usage.get("total_tokens", 0)
def show_cost(self):
print("\n******* OUTPUT SUMMARY *******")
print(f"# of invoke: {self.invoke_count}")
print(f"# of input tokens: {self.input_tokens}")
print(f"# of output tokens: {self.output_tokens}")
print(f"# of total tokens: {self.total_tokens}")
print("******************************")
def main():
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
llm = llm.configurable_fields(max_tokens=ConfigurableField(id="max_tokens"))
prompt = ChatPromptTemplate.from_template("{keyword}とはなんですか？")
chain = prompt | llm | StrOutputParser()
count_tokens_handler = CountTokensHandler()
# invokeの際に、コールバックハンドラを渡すだけ
res = chain.invoke(
{"keyword": "生成AI"}, config={"callbacks": [count_tokens_handler]}
)
print(f"{res=}")
count_tokens_handler.show_cost()
if __name__ == "__main__":
main()
下記が出力結果になります。
$ uv run cost_count_simple.py
res='生成AI（Generative AI）とは、人工知能の一分野であり、データを基に新しいコンテンツを生成する能力を持つAIのことを指します。生成AIは、テキスト、画像、音声、動画など、さまざまな形式のデータを生成することができます。\n\n具体的な例としては、以下のようなものがあります：\n\n1. **テキスト生成**：自然言語処理技術を用いて、文章や詩、ストーリーなどを自動的に生成することができます。例えば、ChatGPTのような対話型AIがこれに該当します。\n\n2. **画像生成**：GAN（Generative Adversarial Networks）やVQ-VAEなどの技術を用いて、新しい画像を生成することができます。例えば、DALL-EやMidjourneyなどのツールが有名です。\n\n3. **音楽生成**：AIが楽曲を作曲したり、特定のスタイルに基づいて音楽を生成することも可能です。\n\n4. **動画生成**：AIを用いて新しい動画コンテンツを作成する技術も進化しています。\n\n生成AIは、クリエイティブな作業の支援や、自動化、データの補完など、さまざまな分野での応用が期待されていますが、同時に著作権や倫理的な問題も議論されています。'
******* OUTPUT SUMMARY *******
# of invoke: 1
# of input tokens: 14
# of output tokens: 330
# of total tokens: 344
******************************ちゃんとOUTPUT SUMMARYに結果が出力されていますね！
呼ばれた回数（# of invok）が1回で、入力されたトークン（# of input tokens）は14、出力トークン（# of output tokens）は330、全トークン（# of total tokens）は334という結果でした。
これなら、割と正確に、OpenAI APIの料金の計算ができそうですね！
これ、かなり便利なのが、LangGraphのGraphでも全く同じコードでトークンを追跡できちゃうんです。
コード全文は書きませんが、下記のような簡単なグラフに対して適用してみます。
簡単に言うと、最初のselectionノードで生成AIの役割を決定させ、answeringノードでその役割に応じて回答させる流れです。
ソースコード（一部抜粋）
workflow = StateGraph(State)
workflow.add_node("selection", selection_node)
workflow.add_node("answering", answering_node)
workflow.set_entry_point("selection")
workflow.add_edge("selection", "answering")
workflow.add_edge("answering", END)
compiled = workflow.compile()
count_tokens_handler = CountTokensHandler()
initial_state = State(query="生成AIについて教えてください")
result = compiled.invoke(
initial_state, config={"callbacks": [count_tokens_handler]}
)
print(result)
count_tokens_handler.show_cost()
実行結果
$ uv run cost_count.py
{'query': '生成AIについて教えてください', 'current_role': '生成AI製品エキスパート', 'messages': ['生成AIとは、人工知能の一分野であり、特にテキスト、画像、音声などのコンテンツを自動的に生成する技術を指します。これらの技術は、機械学習や深層学習を基盤にしており、大量のデータを学習することで新しいコンテンツを創出する能力を持っています。\n\n### 主な特徴と技術\n1. **自然言語処理（NLP）**: テキスト生成において、生成AIは文脈を理解し、自然な言語で文章を作成します。例えば、OpenAIのGPTシリーズやGoogleのBERTなどが代表的です。\n\n2. **画像生成**: DALL-EやMidjourneyなどのモデルは、テキストから画像を生成する能力を持ち、クリエイティブなビジュアルコンテンツを作成します。\n\n3. **音声生成**: 音声合成技術を用いて、テキストを自然な音声に変換することができます。これにより、音声アシスタントやナレーションなどが実現されています。\n\n### 利用例\n- **コンテンツ制作**: ブログ記事、広告コピー、ストーリーなどの自動生成。\n- **デザイン**: グラフィックデザインやプロダクトデザインのアイデア出し。\n- **教育**: 学習教材の作成や個別指導のためのカスタマイズされたコンテンツ提供。\n\n### 課題と倫理\n生成AIには、著作権やフェイクニュースの生成、バイアスの問題など、いくつかの倫理的な課題も存在します。これらの問題に対処するためには、透明性や責任ある使用が求められています。\n\n生成AIは、今後ますます多くの分野で活用されることが期待されており、技術の進化とともに新たな可能性が広がっています。'], 'current_judge': False, 'judgement_reason': ''}
******* OUTPUT SUMMARY *******
# of invoke: 2
# of input tokens: 365
# of output tokens: 461
# of total tokens: 826
******************************まとめ
今回は、簡単に、もとのソースコードを極力汚さないでトークンを追跡する方法を紹介しました。
複数回呼ばれるような場合にも、invokeにcallbacksを指定するだけで追跡してくれるので、とても便利だと思います。
是非使ってみてください。
コメント