Zenn
🔖

Flask/LangChain を使った Chat アプリでチャット履歴をTiDBに保存する

2024/06/08に公開
TiDB
tech

今日はこちらをやっていきます。
https://python.langchain.com/v0.1/docs/integrations/memory/tidb_chat_message_history/
こちらのサンプルではUIがないのでFlaskを使ったUIを最後に作ってみます。

Flask とは

FlaskはPythonで書かれたマイクロウェブフレームワークです。特定のツールやライブラリを必要としないため、マイクロフレームワークに分類されます。 データベース抽象化レイヤー、フォーム検証、または既存のサードパーティライブラリが一般的な機能を提供するその他のコンポーネントはありません。ただし、Flaskは、Flask自体に実装されているかのようにアプリケーション機能を追加できる拡張機能をサポートしています。オブジェクトリレーショナルマッパー、フォーム検証、アップロード処理、さまざまなオープン認証テクノロジー、およびいくつかの一般的なフレームワーク関連ツールの拡張機能が存在します。

(Wikipedia 英語版より）
Pythonベースのコードをウェブ化させるフレームワークで、様々な追加機能を利用可能なツール群が存在しています。

さっそくやってみる

まずは必要なLangChainライブラリをインストールします。

pip install --upgrade --quiet langchain langchain_openai

以下のpythonコードを実行します。

test.py
import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
# copy from tidb cloud console
tidb_connection_string_template = "mysql+pymysql://root:<PASSWORD>@34.212.137.91:4000/test"
tidb_password = getpass.getpass("Input your TiDB password:")
tidb_connection_string = tidb_connection_string_template.replace(
    "<PASSWORD>", tidb_password
)

from datetime import datetime

from langchain_community.chat_message_histories import TiDBChatMessageHistory

history = TiDBChatMessageHistory(
    connection_string=tidb_connection_string,
    session_id="code_gen",
    earliest_time=datetime.utcnow(),  # Optional to set earliest_time to load messages after this time point.
)

history.add_user_message("How's our feature going?")
history.add_ai_message(
    "It's going well. We are working on testing now. It will be released in Feb."
)

print(history.messages)

tidb_connection_string_templateの値は以下の記事を参考に皆さんの値に書き換えてください。
https://zenn.dev/kameping/articles/128ac71b824148

python test.py

そうすると以下の出力が出てきます。

[HumanMessage(content="How's our feature going?"), AIMessage(content="It's going well. We are working on testing now. It will be released in Feb.")]

この時点ではまだOpenAIを使っていません。
単純に人間がHow's our feature going?と質問し、AIがIt's going well. We are working on testing now. It will be released in Feb.と答えた想定で会話がされたというダミーの内容をhistoryに格納し、print(history.messages)で出力しているだけです。
TiDBのマネージメントコンソールではlangchain_message_storeというテーブルが作成され、その中に会話履歴が保存されていることがわかります。


これは以下の部分が処理を抽象化しています。

history = TiDBChatMessageHistory(
    connection_string=tidb_connection_string,
    session_id="code_gen",
    earliest_time=datetime.utcnow(),  # Optional to set earliest_time to load messages after this time point.
)

ポイントとなるのはそれぞれ格納された値が以下であることです。

human
"data": {"additional_kwargs": {}, "content": "How's our feature going?", "example": false, "id": null, "name": null, "response_metadata": {}, "type": "human"}, "type": "human"}

{"data": {"additional_kwargs": {}, "content": "It's going well. We are working on testing now. It will be released in Feb.", "example": false, "id": null, "invalid_tool_calls": [], "name": null, "response_metadata": {}, "tool_calls": [], "type": "ai", "usage_metadata": null}, "type": "ai"}

テーブルに格納された値にはhumanかaiがセットされています。

history.add_user_message("How's our feature going?")
history.add_ai_message(
    "It's going well. We are working on testing now. It will be released in Feb."
)

この部分がそれを処理しています。add_user_messageがhuman、history.add_ai_messageがaiです。

デフォルトではテーブル名はlangchain_message_storeと固定されていますがtable_nameというパラメータを付けると、変更ができます。
https://api.python.langchain.com/en/latest/chat_message_histories/langchain_community.chat_message_histories.tidb.TiDBChatMessageHistory.html

ここからOpenAIとの連携を開始していきます。
先ほどのコードに以下を追記します。

test.py
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You're an assistant who's good at coding. You're helping a startup build",
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)
chain = prompt | ChatOpenAI()

from langchain_core.runnables.history import RunnableWithMessageHistory

chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: TiDBChatMessageHistory(
        session_id=session_id, connection_string=tidb_connection_string
    ),
    input_messages_key="question",
    history_messages_key="history",
)

response = chain_with_history.invoke(
    {"question": "Today is Jan 1st. How many days until our feature is released?"},
    config={"configurable": {"session_id": "code_gen"}},
)
print(response)

以下の出力が追加されています。

content='There are 31 days until Feb 1st when the feature is scheduled to be released.' response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 119, 'total_tokens': 138}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None} id='run-6a89100e-2931-4005-a4e5-7d27ac3946de-0' usage_metadata={'input_tokens': 119, 'output_tokens': 19, 'total_tokens': 138}

では中身を見ていきます。

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You're an assistant who's good at coding. You're helping a startup build",
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)
chain = prompt | ChatOpenAI()

ChatPromptTemplate.from_messages()により対話（プロンプトと言われます)の開始点を宣言しています。つまりOpenAIに対して、あなたは今からどういう振る舞いをする人なんですよ、というのを教えています。例えばその部分を以下に変更すると挙動の違いが判ります。

        (
            "system",
            "You're very bad guy who hates coding. You do not want to help anything",
        ),

content="I'm not sure." response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 228, 'total_tokens': 233}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None} id='run-00940fcf-6893-47ed-97d9-0f9ae6b6de86-0' usage_metadata={'input_tokens': 228, 'output_tokens': 5, 'total_tokens': 233}

MessagesPlaceholder(variable_name="history"),

これによりhistoryに会話履歴が保存されます。

chain = prompt | ChatOpenAI():

定義されたプロンプトをOpenAIに問合せ結果をchainに格納します。history情報がpromptには格納されているため、会話履歴をもとにOpenAIが回答を返すことになります。

chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: TiDBChatMessageHistory(
        session_id=session_id, connection_string=tidb_connection_string
    ),
    input_messages_key="question",
    history_messages_key="history",
)

runnableとはLangChainにとって実行可能な要素、という意味です。先ほど定義したchainを実行可能要素として定義しています。
次のTiDBChatMessageHistoryパートでデータベースへの接続を定義しています。
input_messages_keyは、ユーザーからの入力メッセージを識別するためのキーです。ここでは、質問を識別するために使用されます。history_messages_keyは、過去の会話履歴を識別するためのキーです。ここでは、LangChainの内部で使用される履歴メッセージを指定します。

response = chain_with_history.invoke(
    {"question": "Today is Jan 1st. How many days until our feature is released?"},
    config={"configurable": {"session_id": "code_gen"}},
)

先ほど宣言したRunnable(chain_with_history)をinvokeメソッドで実行しています。

flask を使ったUIの作成と対話型への変更

まずはflaskをインストールします。

pip install flask

つぎにtest.pyを以下に入れ替えます。

test.py
from flask import Flask, request, render_template
import os
from datetime import datetime
from langchain_community.chat_message_histories import TiDBChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain_core.runnables.history import RunnableWithMessageHistory

app = Flask(__name__)

# 環境変数の設定
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
# copy from tidb cloud console
tidb_connection_string_template = "mysql+pymysql://root:<PASSWORD>@34.212.137.91:4000/test"
tidb_password = getpass.getpass("Input your TiDB password:")
tidb_connection_string = tidb_connection_string_template.replace(
    "<PASSWORD>", tidb_password
)

@app.route("/", methods=["GET", "POST"])
def index():
    response_text = ""
    if request.method == "POST":
        user_question = request.form["question"]
        
        # TiDBChatMessageHistoryの初期化
        history = TiDBChatMessageHistory(
            connection_string=tidb_connection_string,
            session_id="code_gen",
            table_name="kamepingchattest",
            earliest_time=datetime.utcnow()
        )

        # ユーザーメッセージを追加
        history.add_user_message(user_question)

        # プロンプトテンプレートの作成
        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    "You're an assistant who's good at coding. You're helping a startup build",
                ),
                MessagesPlaceholder(variable_name="history"),
                ("human", "{question}"),
            ]
        )
        chain = prompt | ChatOpenAI()

        chain_with_history = RunnableWithMessageHistory(
            chain,
            lambda session_id: TiDBChatMessageHistory(
                session_id=session_id, connection_string=tidb_connection_string
            ),
            input_messages_key="question",
            history_messages_key="history",
        )

        # 質問に対する応答を生成
        response = chain_with_history.invoke(
            {"question": user_question},
            config={"configurable": {"session_id": "code_gen"}},
        )

        # 応答の内容を取得
        response_text = response.content

        # 応答を履歴に追加
        history.add_ai_message(response_text)

    return render_template("index.html", response=response_text)

if __name__ == "__main__":
    app.run(debug=True)

次にtemplates/index.htmlを作成します。

templates/index.html
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Conversation</title>
</head>
<body>
    <h1>Conversation</h1>
    <form method="POST">
        <input type="text" name="question" placeholder="Ask something...">
        <input type="submit" value="Ask">
    </form>
    <h2>Response:</h2>
    <p>{{ response }}</p>
    
   
    <script>
        function getConversationHistory() {
            fetch('/conversation_history')
            .then(response => response.json())
            .then(data => {
                // 会話履歴を表示
                console.log(data); // データをコンソールに表示
            });
        }
    </script>
</body>
</html>

実行すると以下が出力されます。

python test.py   
 * Serving Flask app 'test'
 * Debug mode: on
WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
 * Running on http://127.0.0.1:5000
Press CTRL+C to quit
 * Restarting with stat
 * Debugger is active!
 * Debugger PIN: 942-824-857

http://127.0.0.1:5000にアクセスします。以下のようなChatアプリが出来上がります。

TiDB側では以下のように会話履歴が蓄積されています。

human
{"data": {"additional_kwargs": {}, "content": "こんにちは！", "example": false, "id": null, "name": null, "response_metadata": {}, "type": "human"}, "type": "human"}
ai
{"data": {"additional_kwargs": {}, "content": "こんにちは！お困り事はありますか？", "example": false, "id": "run-942c69e6-da7a-4189-9499-58dc4cc9b6bc-0", "invalid_tool_calls": [], "name": null, "response_metadata": {"finish_reason": "stop", "logprobs": null, "model_name": "gpt-3.5-turbo", "system_fingerprint": null, "token_usage": {"completion_tokens": 12, "prompt_tokens": 243, "total_tokens": 255}}, "tool_calls": [], "type": "ai", "usage_metadata": {"input_tokens": 243, "output_tokens": 12, "total_tokens": 255}}, "type": "ai"}

Discussion