try:
    import unzip_requirements
except ImportError:
    pass

import json
from langchain import OpenAI
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.faiss import FAISS

chain = load_qa_with_sources_chain(OpenAI(temperature=0))

search_index = FAISS.load_local(".", OpenAIEmbeddings())


def print_answer(query):
    print(query)
    input_documents = search_index.similarity_search(query, k=3)
    for i in range(0, 3):
        print(f"候補{i+1}: {input_documents[i].metadata}")

    result = \
        chain(
            {
                "input_documents": input_documents,
                "question": query,
            }, return_only_outputs=True)
    answer = result["output_text"]
    print(answer)
    return answer


def hello(event, context):
    query = json.loads(event["body"])["text"].strip()

    return {
        "statusCode": 200,
        "body": print_answer(query)
    }