TiDB

TiDB Cloud 是一個全面的資料庫即服務 (DBaaS) 解決方案，提供專用和無伺服器選項。TiDB Serverless 現在將內建的向量搜尋整合到 MySQL 環境中。透過此增強功能，您可以使用 TiDB Serverless 無縫開發 AI 應用程式，而無需新的資料庫或其他技術堆疊。建立免費的 TiDB Serverless 叢集，並在 https://pingcap.com/ai 開始使用向量搜尋功能。

本筆記本介紹如何使用 TiDB 儲存聊天訊息歷史記錄。

設定

首先，我們將安裝以下相依性

%pip install --upgrade --quiet langchain langchain_openai langchain-community

設定您的 OpenAI 金鑰

import getpass
import os

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass.getpass("Input your OpenAI API key:")

最後，我們將設定與 TiDB 的連線。在本筆記本中，我們將遵循 TiDB Cloud 提供的標準連線方法，以建立安全且有效率的資料庫連線。

# copy from tidb cloud console
tidb_connection_string_template = "mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
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."
)

API 參考：TiDBChatMessageHistory

history.messages

[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.")]

與歷史資料聊天

讓我們以先前產生的歷史資料為基礎，建立動態的聊天互動。

首先，使用 LangChain 建立聊天鏈

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

API 參考：ChatPromptTemplate | MessagesPlaceholder | ChatOpenAI

在歷史記錄上建立 Runnable

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",
)

API 參考：RunnableWithMessageHistory

啟動聊天

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

AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')

檢查歷史資料

history.reload_cache()
history.messages

[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."),
 HumanMessage(content='Today is Jan 1st. How many days until our feature is released?'),
 AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')]