Lambda+Bedrock+TiDB Cloudで、LangChainを使って会話履歴を含めたやり取りをしてみた

CREATE TABLE chat_history (
    id BIGINT AUTO_INCREMENT PRIMARY KEY,
    session_id VARCHAR(255) NOT NULL,
    message TEXT,
    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
);

import json
import os
from langchain_community.chat_message_histories.tidb import TiDBChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, AIMessage, SystemMessage
from langchain_aws.chat_models.bedrock import ChatBedrock
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda, ConfigurableFieldSpec, RunnablePassthrough
from langchain_community.chat_message_histories import ChatMessageHistory

# DB接続用の文字列を作成
connection_string = (
    f"mysql+pymysql://{os.getenv('TIDB_USER')}:{os.getenv('TIDB_PASSWORD')}@"
    f"{os.getenv('TIDB_HOST')}:{os.getenv('TIDB_PORT')}/{os.getenv('TIDB_DATABASE')}?"
    f"ssl_ca=./isrgrootx1.pem"
)

# 履歴格納用
store = {}

# セッションIDごとの履歴の取得
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = TiDBChatMessageHistory(
            table_name="chat_history",
            connection_string=connection_string,
            session_id=session_id
        )
    return store[session_id]

# プロンプトテンプレートの作成
# SystemMessageの部分は、("system", "指示内容")でも良いと思います
prompt_template = ChatPromptTemplate.from_messages(
    [
        SystemMessage(content="あなたは優秀なAIです。あなたは昨日シチューを食べました。設定上、あなたは食事をするAIということになっています。"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)

# モデル指定（Claude 3.5 Sonnet）
model = ChatBedrock(
    model_id="anthropic.claude-3-5-sonnet-20240620-v1:0",
    region_name="ap-northeast-1",
)

# Runnableの作成
runnable = prompt_template | model

chain_with_history = RunnableWithMessageHistory(
    runnable=runnable,
    get_session_history=get_session_history,
    input_messages_key="question",
    history_messages_key="history",
)

# Lambdaメイン
def lambda_handler(event, context):
    user_input = event.get("input")
    session_id = event.get("sessionId")

    print("User input:", user_input)

    # AIの応答を生成
    response = chain_with_history.invoke(
        {"question": user_input},
        config={"configurable": {"session_id": session_id}},
    )

    print("Response:", response)

    # レスポンスを返す
    return {
        "statusCode": 200,
        "body": json.dumps({"response": response.content})
    }

langchain
langchain-aws
langchain-core
langchain-community
pymysql
pydantic

pip install -r requirements.txt -t package/

{
    "input": "昨日は何を食べましたか？",
    "sessionId": "hugahuga"
}

{
    "input": "私は先ほどした質問は何ですか？",
    "sessionId": "hugahuga"
}