RAGって一体どれくらいトークンを消費するの？

ところで、RAGってどれくらいトークンを消費するのでしょうか？
シンプルな構成のRAGを作って検証しました。

トークン数の取得方法

LangChainにはget_bedrock_anthropic_callbackという便利なものが提供されています。

こんな感じで使います。

with get_bedrock_anthropic_callback() as cb:
    result = llm.invoke("Tell me a joke")
    result2 = llm.invoke("Tell me a joke")
    print(cb)

OpenAI用にはget_openai_callbackというものがあります。

構成

LangChainのチュートリアルを参考にしました。

日本語で検証したかったので、こちらのサイトの内容をRAGに取り込みました。

Retriever

WebBaseLoaderで読み込んだデータをチャンク分割し、Chromaに登録します。

チャンク分割はRecursiveCharacterTextSplitterを使い、チャンクサイズを1000、チャンクオーバーラップを200としています。

loader = WebBaseLoader(
    web_paths=(
        "https://www.aboutamazon.jp/news/aws/new-additions-to-amazon-bedrock-make-it-easier-and-faster-than-ever-to-build-generative-ai-applications-securely",
    ),
    bs_kwargs=dict(parse_only=bs4.SoupStrainer(class_=("ArticlePage-main"))),
)
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
vectorstore = Chroma.from_documents(
    documents=splits,
    embedding=BedrockEmbeddings(model_id="amazon.titan-embed-text-v2:0"),
)

retriever = vectorstore.as_retriever()

as_retriever()は、取得件数が未指定だとドキュメントを4件取得します。

Prompt

LangChain Hubからrlm/rag-promptを取得し利用します。（こちら）

prompt = hub.pull("rlm/rag-prompt")

プロンプトの内容はこのようなものでした。

You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Question: {question}
Context: {context}
Answer:

LLM

Claude 3 Haikuを使用します。

llm = ChatBedrock(model_id="anthropic.claude-3-haiku-20240307-v1:0")

Chain

組み合わせてChainを作ります。

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

呼び出し部分はこのようになります。

with get_bedrock_anthropic_callback() as cb:
    result = rag_chain.invoke("Titanに新しく追加になったモデルについて教えて下さい。")
    print("---")
    print(result)
    print("---")
    print(cb)

やってみた

さて、結果発表です。

---
質問:Titanに新しく追加になったモデルについて教えて下さい。

答え:
1つ目の新しいモデルはAmazon Titan Image Generatorで、文章を入力するだけで美しい画像を生成できます。
2つ目の新しいモデルはAmazon Titan Text Embeddings V2で、低遅延のモバイル展開から高精度の非同期ワークフローまで、柔軟な埋め込みサイズを提供しています。
---
Tokens Used: 2351
Prompt Tokens: 2202
Completion Tokens: 149
Successful Requests: 1
Total Cost (USD): $0.00073675

ということで、

Inputトークンが2200ぐらい、Outputトークンが150ぐらい 消費します。

※1000トークンでチャンク分割したドキュメントを4件取得の場合

ドキュメント取得を8件にした場合

今回対象にしたドキュメントは、全部で8チャンクに分かれていましたので、8チャンク全部渡した場合はどうなるでしょう。

as_retrieverで取得件数を指定します。

retriever = vectorstore.as_retriever(search_kwargs={"k": 8})

さて、結果は

Tokens Used: 4418
Prompt Tokens: 4139
Completion Tokens: 279
Successful Requests: 1
Total Cost (USD): $0.0013835000000000002

はい。当たり前のように増えました。

与えるドキュメントが増えるとその分饒舌になり、Outputトークンも増えるのかもしれません。

Sonnet、Opusの場合は？

ドキュメント取得を4件に戻し、モデルを切り替えて実施しました。

モデル	Inputトークン	Outputトークン	価格（$）	価格（円：$1.00=150円）
Haiku	2202	147	$0.00073425	0.1101375
Sonnet	2202	136	$0.008646000000000001	1.29690000000000015
Opus	-	-	-	-

OpusはToo many requestsエラーになったので断念。Haikuと同じトークン数だったとするとこちら。

モデル	Inputトークン	Outputトークン	価格（$）	価格（円：$1.00=150円）
Opus	2202	147	$0.044055	6.60825

Haiku安っ！！

おまけ

検証ソースコード

import bs4
from langchain import hub
from langchain_aws import BedrockEmbeddings, ChatBedrock
from langchain_chroma import Chroma
from langchain_community.callbacks.manager import get_bedrock_anthropic_callback
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter


# Load, chunk and index the contents of the blog.
loader = WebBaseLoader(
    web_paths=(
        "https://www.aboutamazon.jp/news/aws/new-additions-to-amazon-bedrock-make-it-easier-and-faster-than-ever-to-build-generative-ai-applications-securely",
    ),
    bs_kwargs=dict(parse_only=bs4.SoupStrainer(class_=("ArticlePage-main"))),
)
docs = loader.load()


text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
vectorstore = Chroma.from_documents(
    documents=splits,
    embedding=BedrockEmbeddings(model_id="amazon.titan-embed-text-v2:0"),
)

# Retrieve and generate using the relevant snippets of the blog.
retriever = vectorstore.as_retriever()
prompt = hub.pull("rlm/rag-prompt")


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


model_ids = [
    "anthropic.claude-3-haiku-20240307-v1:0",
    "anthropic.claude-3-sonnet-20240229-v1:0",
    "anthropic.claude-3-opus-20240229-v1:0",
]

for model_id in model_ids:

    llm = ChatBedrock(model_id=model_id, region_name="us-west-2")

    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    with get_bedrock_anthropic_callback() as cb:
        result = rag_chain.invoke(
            "Titanに新しく追加になったモデルについて教えて下さい。"
        )
        print("---")
        print(model_id)
        print("---")
        print(result)
        print("---")
        print(cb)

---

bedrock_anthropic_callbackのソースコードを確認すると、llm_outputの値を取得して算出していました。StrOutputParser()をコメントアウトしてllmのinvoke結果を取得すると消費プロンプト数が取得できました。

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    # | StrOutputParser()
)

result = rag_chain.invoke("Titanに新しく追加になったモデルについて教えて下さい。")
print(json.dumps(result.response_metadata, indent=2))

{
  "usage": {
    "prompt_tokens": 2202,
    "completion_tokens": 97,
    "total_tokens": 2299
  },
  "stop_reason": "end_turn",
  "model_id": "anthropic.claude-3-haiku-20240307-v1:0"
}