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

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

## トークン数の取得方法

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

https://python.langchain.com/v0.2/docs/how_to/chat_token_usage_tracking/#bedrock-anthropic

こんな感じで使います。

```python
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のチュートリアルを参考にしました。

https://python.langchain.com/v0.2/docs/tutorials/rag/

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

https://www.aboutamazon.jp/news/aws/new-additions-to-amazon-bedrock-make-it-easier-and-faster-than-ever-to-build-generative-ai-applications-securely

### Retriever

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

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

```python
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()
```

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

### Prompt

LangChain Hubから`rlm/rag-prompt`を取得し利用します。（[こちら](https://smith.langchain.com/hub/rlm/rag-prompt)）

```python
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を使用します。

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

### Chain

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

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

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

```python
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`で取得件数を指定します。

```python
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

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

:::note
与えるドキュメントが増えるとその分饒舌になり、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 |

:::note
Haiku安っ！！
:::

## おまけ


```python:検証ソースコード
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のソースコード](https://github.com/langchain-ai/langchain/blob/master/libs/community/langchain_community/callbacks/bedrock_anthropic_callback.py)を確認すると、llm_outputの値を取得して算出していました。`StrOutputParser()`をコメントアウトしてllmの`invoke`結果を取得すると消費プロンプト数が取得できました。

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

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

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