import langchain
from dotenv import load_dotenv
from langchain.chains import FlareChain
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.vectorstores import FAISS

langchain.verbose = True

load_dotenv()

# FAISSで保存されたベクトルを読み込む
embeddings = OpenAIEmbeddings()
db = FAISS.load_local("./tmp/faiss", embeddings)
retriever = db.as_retriever()

# FLAREの準備
chat = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
flare = FlareChain.from_llm(
    llm=chat,
    retriever=retriever,
    min_prob=0.2,
)
# 回答の生成には、logprobsが使えるCompletions APIを使う
flare.response_chain.llm = OpenAI(
    model="gpt-3.5-turbo-instruct",
    temperature=0,
    model_kwargs={"logprobs": 1},
)

query = "LangChainとは"
result = flare.run(query)
print(result)

langchain.vectorstores.FAISS.load_local langchain.llms.OpenAI langchain.chat_models.ChatOpenAI langchain.chains.FlareChain.from_llm langchain.embeddings.OpenAIEmbeddings
[((285, 298), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (296, 298), False, 'from dotenv import load_dotenv\n'), ((336, 354), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (352, 354), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((360, 403), 'langchain.vectorstores.FAISS.load_local', 'FAISS.load_local', (['"""./tmp/faiss"""', 'embeddings'], {}), "('./tmp/faiss', embeddings)\n", (376, 403), False, 'from langchain.vectorstores import FAISS\n'), ((453, 506), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)'}), "(model_name='gpt-3.5-turbo', temperature=0)\n", (463, 506), False, 'from langchain.chat_models import ChatOpenAI\n'), ((515, 579), 'langchain.chains.FlareChain.from_llm', 'FlareChain.from_llm', ([], {'llm': 'chat', 'retriever': 'retriever', 'min_prob': '(0.2)'}), '(llm=chat, retriever=retriever, min_prob=0.2)\n', (534, 579), False, 'from langchain.chains import FlareChain\n'), ((663, 751), 'langchain.llms.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo-instruct"""', 'temperature': '(0)', 'model_kwargs': "{'logprobs': 1}"}), "(model='gpt-3.5-turbo-instruct', temperature=0, model_kwargs={\n    'logprobs': 1})\n", (669, 751), False, 'from langchain.llms import OpenAI\n')]