LangChainでChatGPTに検索とDBの機能を追加する

2023年4月21日 10:14

はじめに

人には、数時間の間保管される短期記憶と年単位で保管される長期記憶があり、短期記憶は、数秒から数時間の記憶で、リアルタイムでの対話やタスクを遂行する際に役立ちます。一方、長期記憶は、年単位で大量の情報を保管し、過去の経験や学習を現在に活かします。これをAIに当てはめるとPromptを通したAIとの会話は、短期記憶に相当するのではないでしょうか。

一方、長期記憶に相当する機能はChatGPTなどのAI単体では実現できないためベクトルデータベースを利用することが必要になります。

また、Google検索を可能とすることで、AIは最新の情報や専門知識の利用も可能になります。

Google Colabで実行

環境変数にOpenAIとGoogle Search APIのトークンを設定

# OpneAIとGoogle Search APIキー
%env OPENAI_API_KEY = sk-xxxxxxxxxxxxxxxxxxx
%env GOOGLE_CSE_ID = xxxxxxxxxxxxxxxxxxxxxxx
%env GOOGLE_API_KEY = xxxxxxxxxxxxxxxxxxxxxx

パッケージのインストール

!pip install langchain
!pip install openai
!pip install gradio
!pip install chromadb
!pip install tiktoken
!pip install gradio
!pip install unstructured

Vector DB を作るまでに必要なパッケージをインポート

# テキストファイルの読み込み
from langchain.document_loaders import UnstructuredFileLoader
# from langchain.text_splitter import CharacterTextSplitter
from langchain.text_splitter import RecursiveCharacterTextSplitter
# Embedding
from langchain.embeddings.openai import OpenAIEmbeddings
# vectorDB
from langchain.vectorstores import Chroma

テキストファイルを読み込んでVector DBを作成

今回はジョジョの奇妙な冒険 Part3 スターダストクルセイダースのwikiを利用しました。
テキストファイル以外にもより実用的なソースとして例えば、Notion DB なども利用可能です。
大規模なデータを利用する際はPineconeなど外部のVectorDBのサービスを利用することになるかと思います。

# テキストファイルの読み込みとチャンクへ分割
loader = UnstructuredFileLoader("jojo3.txt")
documents = loader.load()
# text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# Embedding
embeddings = OpenAIEmbeddings()
# vectorDBの作成
db = Chroma.from_documents(docs, embeddings)

Tool、Agent、GUIのパッケージをインポート

# Google Search API
from langchain.utilities import GoogleSearchAPIWrapper
# Memory
from langchain.memory import ConversationBufferMemory
# ChatOpenAI
from langchain.chat_models import ChatOpenAI
# Retriever
from langchain.chains import RetrievalQA
# agentとtool
from langchain.agents import Tool, initialize_agent
# prompt
from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
from langchain import LLMChain
# UI
import gradio as gr

AI に使用させるツール (Google検索とVectorDB) を定義

# GoogleSearchAPIとvectorDB
search = GoogleSearchAPIWrapper()

jojo3 = RetrievalQA.from_chain_type(llm=ChatOpenAI(), chain_type="map_reduce", retriever=db.as_retriever())
tools = [
    Tool(
        name = "Search",
        func=search.run,
        description="useful for when you need to answer questions about current events"
    ),
    Tool(
        name = "database of jojo part3",
        func=jojo3.run,
        description="useful for when you need to answer question about jojo part3."
    )
]
# agent が使用する memory
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

Promptを定義

prefixではツールを使えるとことを伝え、suffixでは承太郎っぽく回答してくれるように語尾を指定

prefix = """次の質問にできる限り日本語で答えてください。次のツールにアクセスできます:"""
suffix = """始めましょう! 最終的な答えを出すときは、回答の後に"やれやれだぜ"を追加してください

Question: {input}
{agent_scratchpad}"""

prompt = ZeroShotAgent.create_prompt(
    tools,
    prefix=prefix,
    suffix=suffix,
    input_variables=["input", "agent_scratchpad"]
)

Agentを定義

promptとLLMをLLMChainで組み合わせる

llm_chain = LLMChain(llm=ChatOpenAI(temperature=0), prompt=prompt)

# エージェント
agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools)

# AgentExecuterで実行
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory)

GUIの定義

responseでagent_executorを使用することを指定

def chat(message, history):
  history  = history or []
  response = agent_executor.run(input = message) 
  
  history.append((message, response))
  
  return history, history

chatbot = gr.Chatbot().style(color_map=('green', 'pink'))

demo = gr.Interface(
    chat,
    ['text',  'state'],
    [chatbot, 'state'],
    allow_flagging = 'never',
)

会話の開始

demo.launch()