langchainとDatabricksで(私が)学ぶAgent : ローカルLLMでReAct

2024年はLLM Agent飛躍の年です。たぶん。2023年からそうだというツッコミは無しで。

# 導入

少し前に[Langchainのバージョン0.1.0がリリース](https://blog.langchain.dev/langchain-v0-1-0/)されました。
まだまだ雑多感があるとはいえ、モジュールもかなり整理され、LCELをベースとした記法で見通しもだいぶよくなってきました。（Productionに採用するにあたっては多くの苦労が待ち受けていそうですが）

Agent関連の機能もだいぶ整理されたり、[langgraph](https://python.langchain.com/docs/langgraph)が登場したこともあり、最近の記法（と最近のLLMを使って）Agent周りを今更ながらちゃんとやってみたくなりました。

まずはローカルLLMでReActを実践してみます。
今回は、シンプルにWeb検索を行って質問結果を返すAgentを実装しました。

LangchainのAgentsとは何？という方は[公式ドキュメント](https://python.langchain.com/docs/modules/agents/)を是非一読ください。英語ですが、だいぶ内容が整理されて読みやすいくなってきてるなと思います。
ReActについては、以下のサイトが参考になるかと思います。

https://book.st-hakky.com/data-science/llm-prompt-engineering-react/

実践・検証はDatabricks on AWS上で実施しました。

# Step1. パッケージのインストール

使うパッケージをインストール。
今回も[ExLlama V2](https://github.com/turboderp/exllamav2)を使って推論します。

```python
%pip install -U transformers accelerate "exllamav2>=0.0.11" langchain langchainhub duckduckgo-search

dbutils.library.restartPython()
```

# Step2. モデルのロード

LLMのモデルを読み込みます。
今回も[こちら](https://qiita.com/isanakamishiro2/items/c1856c94d4c0227a0d9e)で作成した、langchainのカスタムチャットモデルを使って読み込みました。

使っているモデルは[OpenChat 3.5の1210版](https://huggingface.co/openchat/openchat-3.5-1210)をGPTQで量子化した以下のモデルです。
※ OpenChat 3.5は0106版という新しいバージョンが出ているのですが、思ったような応答にならなかったため、今回は1210版にしています。

https://huggingface.co/TheBloke/openchat-3.5-1210-GPTQ



```python
from exllamav2_chat import ChatExllamaV2Model

model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--openchat-3.5-1210-GPTQ"
chat_model = ChatExllamaV2Model.from_model_dir(
    model_path,
    cache_max_seq_len=8192,
    system_message_template="GPT4 Correct User: {}
GPT4 Correct Assistant: OK.
",    
    human_message_template="GPT4 Correct User: {}
GPT4 Correct Assistant: ",
    ai_message_template="{}",
    temperature=0.0001,
    top_p=0.0001,
    max_new_tokens=1024,
    repetition_penalty = 1.15,
    low_memory=True,
    cache_8bit=True,
)
```

# Step3. ツールの準備

LLM Agentの中で呼び出して実行するツールを定義します。

今回は[DuckDuckGo](https://duckduckgo.com/)を使ってWeb検索するツールを使うことにします。
余談ですが、[Tavily](https://tavily.com/)というLLMからの利用に最適化したWeb検索APIサービスもあるんですね。。。langchainのドキュメントを読んでいて初めて知りました。

```python
from langchain_community.tools.ddg_search.tool import DuckDuckGoSearchResults, DuckDuckGoSearchRun

search = DuckDuckGoSearchRun(max_results=1)

# 動作確認
# search.run("what is the weather in SF")

# ↓のような内容が返ってくる
# "San Francisco Weather Forecasts. Weather Underground provides local & long-range weather forecasts, weatherreports, maps & tropical weather conditions ...
```

# Step4. プロンプトの準備

[LangSmith hub](https://smith.langchain.com/hub)上で公開されているReAct用のプロンプトを取得します。
今回はChatModelに対応し、かつ結果をJSONで返す以下のプロンプトテンプレートを使用しました。

https://smith.langchain.com/hub/hwchase17/react-chat-json

```python
from langchain import hub
prompt = hub.pull("hwchase17/react-chat-json")
```

なお、以下のような内容になります。
ReActのプロンプトはそれなりに複雑なので、簡単にプロンプトテンプレートが取得できるhub機能は便利ですね。

```python
ChatPromptTemplate(
    input_variables=["agent_scratchpad", "input", "tool_names", "tools"],
    input_types={
        "chat_history": typing.List[
            typing.Union[
                langchain_core.messages.ai.AIMessage,
                langchain_core.messages.human.HumanMessage,
                langchain_core.messages.chat.ChatMessage,
                langchain_core.messages.system.SystemMessage,
                langchain_core.messages.function.FunctionMessage,
                langchain_core.messages.tool.ToolMessage,
            ]
        ],
        "agent_scratchpad": typing.List[
            typing.Union[
                langchain_core.messages.ai.AIMessage,
                langchain_core.messages.human.HumanMessage,
                langchain_core.messages.chat.ChatMessage,
                langchain_core.messages.system.SystemMessage,
                langchain_core.messages.function.FunctionMessage,
                langchain_core.messages.tool.ToolMessage,
            ]
        ],
    },
    messages=[
        SystemMessagePromptTemplate(
            prompt=PromptTemplate(
                input_variables=[],
                template="Assistant is a large language model trained by OpenAI.\n\nAssistant is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n\nAssistant is constantly learning and improving, and its capabilities are constantly evolving. It is able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. Additionally, Assistant is able to generate its own text based on the input it receives, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.\n\nOverall, Assistant is a powerful system that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.",
            )
        ),
        MessagesPlaceholder(variable_name="chat_history", optional=True),
        HumanMessagePromptTemplate(
            prompt=PromptTemplate(
                input_variables=["input", "tool_names", "tools"],
                template='TOOLS\n------\nAssistant can ask the user to use tools to look up information that may be helpful in answering the users original question. The tools the human can use are:\n\n{tools}\n\nRESPONSE FORMAT INSTRUCTIONS\n----------------------------\n\nWhen responding to me, please output a response in one of two formats:\n\n**Option 1:**\nUse this if you want the human to use a tool.\nMarkdown code snippet formatted in the following schema:\n\n```json\n{{\n    "action": string, \\ The action to take. Must be one of {tool_names}\n    "action_input": string \\ The input to the action\n}}\n```\n\n**Option #2:**\nUse this if you want to respond directly to the human. Markdown code snippet formatted in the following schema:\n\n```json\n{{\n    "action": "Final Answer",\n    "action_input": string \\ You should put what you want to return to use here\n}}\n```\n\nUSER\'S INPUT\n--------------------\nHere is the user\'s input (remember to respond with a markdown code snippet of a json blob with a single action, and NOTHING else):\n\n{input}',
            )
        ),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ],
)
```

# Step5. Agent用Chainの作成

Agent用のChainを作成し、実行のためのAgentExecutorを作成します。
ツールとして、Step3で作成したDuckDuckGo検索用ツールを渡しています。
これによって、指定されたクエリに対しWeb検索を行い、その結果を用いてクエリの回答を得るというエージェントが出来上がります。

```python
from langchain.agents import AgentExecutor, create_json_chat_agent

tools = [search]

# Construct the ReAct agent
agent = create_json_chat_agent(chat_model, tools, prompt)

# Create an agent executor by passing in the agent and tools
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
```

なお、```create_json_chat_agent```関数で得られるオブジェクトは、内部的にはLCELで記述されたチェーン(Runnableオブジェクト)です。

```python
type(agent)
# langchain_core.runnables.base.RunnableSequence
```

これでAgentを実行できる準備が整いました。

# Step6. 実行してみる

いくつか簡単なクエリを実行してみます。
まずは時事系の質問。

```python
agent_executor.invoke({"input": "日本の総理大臣は誰？"})
```

````:出力

> Entering new AgentExecutor chain...

```json


> Entering new AgentExecutor chain...

```json
{
    "action": "duckduckgo_search",
    "action_input": "who is the prime minister of japan"
}
```KAMIKAWA Yoko. House of Representatives. Minister of Finance. Minister of State for Financial Services. Minister in charge of Overcoming Deflation. SUZUKI Shunichi. House of Representatives. Minister of Education, Culture, Sports, Science and Technology. MORIYAMA Masahito. Suga Yoshihide See all media Category: History & Society Born: December 6, 1948, Ogachi (now in Yuzawa), Akita prefecture, Japan (age 75) Title / Office: prime minister (2020-2021), Japan Political Affiliation: Liberal-Democratic Party of Japan See all related content → Prime Minister's Office of Japan. The Prime Minister's Major Initiatives. Press # to narrow down the list to related items. #WageIncrease; #NewFormCapitalism ... Instructions by the Prime Minister Regarding the Collision of JCG JA722A and JAL JL516 (18:05) January 2, 2024. Japan's Prime Minister Kishida Fumio, front center, and his new Cabinet pose for a photo at the prime minister's office in Tokyo, Sept. 13, 2023. Credit: Kazuhiro Nogi/Pool Photo via AP. Japan's Prime Minister Fumio Kishida said he was "speechless" by the severity and immense destruction of the disaster zone he saw during the helicopter ride on his first visit to the country's north-central region of Noto since the deadly Jan. 1 earthquakes, amid worries about spreading diseases in evacuation centers.

```json
{
    "action": "Final Answer",
    "action_input": "The current Prime Minister of Japan is Fumio Kishida."
}
```

> Finished chain.
{'input': '日本の総理大臣は誰？',
 'output': 'The current Prime Minister of Japan is Fumio Kishida.'}
````

DuckDuckGoからは正直よくわからない回答が返ってきてるのですが、結果的に正しい回答が返ってきました。
ただ、コンテキストから回答したのかどうかは、ちょっと怪しい気がします。


Web検索しなくても知ってそうなことを聞いてみます。

```python
agent_executor.invoke({"input": "データブリックスとは何ですか？"})
```

````:出力

> Entering new AgentExecutor chain...

```json
{
    "action": "Final Answer",
    "action_input": "データブリックス（DataBricks）は、大規模なデータ処理や機械学習に特化したプラットフォームであり、PythonやRなどのプログラミング言語を使用してデータ分析および機械学習モデルを構築・実行することが可能です。また、データブリックスは、クラウド上で動作するSparkクラスターを提供し、高速なデータ処理と大量のメモリを提供しています。"
}
```

> Finished chain.
{'input': 'データブリックスとは何ですか？',
 'output': 'データブリックス（DataBricks）は、大規模なデータ処理や機械学習に特化したプラットフォームであり、PythonやRなどのプログラミング言語を使用してデータ分析および機械学習モデルを構築・実行することが可能です。また、データブリックスは、クラウド上で動作するSparkクラスターを提供し、高速なデータ処理と大量のメモリを提供しています。'}
````
何度か繰り返してみましたが、今回のモデル&プロンプトだとこのような内容はツールを使わずに回答されてしまいました。

では、次によくある天気を聞いてみましょう。

```python
agent_executor.invoke({"input": "明日の東京の天気は何？"})
```

````:出力

> Entering new AgentExecutor chain...

```json
{
    "action": "duckduckgo_search",
    "action_input": "明日の東京の天気は何？"
}
```東京都の天気 13(土) 14(日) 15(月) 16(火) 17(水) 18(木) 19(金) 20(土) 21(日) 22(月) 23(火) 13日18:00発表 01月13日 ( 土 ) 千代田区 12/6 50% 新宿区 11/5 50% 世田谷区 12/4 50% 三鷹市 12/2 50% 八王子市 12/-2 40% 青梅市 10/-2 40% 大島町 13/7 90%... 1月15日(月) 、東京都の天気をお伝えします。 ... 東京都の明日の天気. ... 風も強く、防寒が必須の一日に。厚手のコートや冬小物を活用して、体感温度の調節をしてください。海上は波が高くなります。 ... 21:00 22:00 23:00 薄い雲 晴れ 晴れ 2°C 2°C 2°C -5°C -5°C -4°C 1月13日（土） 1月14日（日） 都道府県 東京都/東京 の天気 2024年01月13日17時00分発表 都市を選ぶ 東京 大島 八丈島 父島 関東甲信越地方の天気へ 全国の天気へ 週間天気 信頼度について 3日目以降の降水の有無の予報について「予報が適中しやすい」ことと「予報が変わりにくい」ことを表す情報でA、B、Cの3段階で表します。... 10℃. 1℃. 東京都東京地方 東京の天気予報です。. 東京の今日・明日の天気予報、3時間毎の天気予報、週間天気予報など素早くチェックできます。. さらに詳細地域を選択してピンポイントな地域の天気予報も提供しています。.

```json
{
    "action": "Final Answer",
    "action_input": "The weather forecast for Tokyo tomorrow is cold with possible snowfall."
}
```

> Finished chain.
{'input': '明日の東京の天気は何？',
 'output': 'The weather forecast for Tokyo tomorrow is cold with possible snowfall.'}
````

"明日"が何日か不透明なので、実際のところ正しくない回答が得られているように見えます。
とはいえ、天気を検索して、その回答を基にきちんと回答する流れは組み立てられていました。

# まとめ

ローカルLLMを使ってReActのロジックで動くAgentを組み立ててみました。

LangchainのドキュメントでもAgentはOpenAIやAWS Bedrockを使う例しか見当たらなかったのですが、ローカルLLMでも（また、7Bという少ないパラメータでも）単純なものは動作することが確認できました。
Web検索するエージェントはよくあるユースケースだと思うのですが、使う検索サービスを[Tavily](https://tavily.com/)に変えたり、またプロンプトをモデルに合わせて適切に変更すれば、より精度よいエージェントになると思います。回答を常に日本語にするとか。

今後、もう少し複雑かつ実用的なエージェントや、LCELでのAgent記法、また[LangGraph](https://python.langchain.com/docs/langgraph)を使ったAgentの構築など、いろいろ試してみたいと思います。