LangChain で GPT-3.5 に Google 検索と Vector Index を使い分けて回答させる【Agent】

# Google Search API
from langchain.utilities import GoogleSearchAPIWrapper
# Memory
from langchain.memory import ConversationBufferMemory
# ChatOpenAI GPT 3.5
from langchain.chat_models import ChatOpenAI
# Embedding用
from langchain.embeddings.openai import OpenAIEmbeddings
# Vector 格納 / FAISS
from langchain.vectorstores import FAISS
# テキストファイルを読み込む
# Q&A用DBQA
from langchain import VectorDBQA

# agentと agentが使用するtool
from langchain.agents import Tool, initialize_agent

# env に読み込ませるAPIキーの類
import key

# 環境変数にAPIキーを設定
import os
os.environ["OPENAI_API_KEY"] = key.OPEN_API_KEY
os.environ["GOOGLE_CSE_ID"] = key.GOOGLE_CSE_ID
os.environ["GOOGLE_API_KEY"] = key.GOOGLE_API_KE

# agent が使用するGoogleSearchAPIWrapperのツールを作成
search = GoogleSearchAPIWrapper()

# agent が使用する vector index
embeddings = OpenAIEmbeddings()
db = FAISS.load_local("faiss_index", embeddings)

jirei = VectorDBQA.from_chain_type(llm=ChatOpenAI(), chain_type="map_reduce", vectorstore=db)

tools = [
    Tool(
        name = "Search",
        func=search.run,
        description="useful for when you need to answer questions about current events"
    ),
    Tool(
        name = "Report of A corp.",
        func=jirei.run,
        description="useful for when you need to answer question about A corp."
    )
]
# agent が使用する memory の作成
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# Agent用 prefix, suffix
prefix = """Anser the following questions as best you can, but speaking Japanese. You have access to the following tools:"""
suffix = """Begin! Remember to speak Japanese when giving your final answer."""

# agent の使用する LLM
llm=ChatOpenAI(temperature=0)

# tool, memory, llm を設定して agent を作成
agent_chain = initialize_agent(tools, llm, agent="chat-conversational-react-description", verbose=True, memory=memory, prefix=prefix, suffix=suffix)

# 会話ループ
user = ""
while user != "exit":
    user = input("何か質問してください。")
    print(user)
    # 入力はagent経由
    ai = agent_chain.run(input=user)
    print(ai)