!pip install llama-index
!pip install langchain

from llama_index import download_loader

CJKPDFReader = download_loader("CJKPDFReader")

loader = CJKPDFReader()
documents = loader.load_data(file="data/aws-parallelcluster/aws-parallelcluster-ug.pdf")

from llama_index import GPTSimpleVectorIndex

index = GPTSimpleVectorIndex.from_documents(documents)

index.save_to_disk('data/aws-parallelcluster/aws-parallelcluster.json')

import gradio as gr
from llama_index import (
    LLMPredictor,
    ServiceContext,
    GPTSimpleVectorIndex,
    QuestionAnswerPrompt,
    RefinePrompt,
)
from langchain.chat_models import ChatOpenAI
import os
import glob

llm_predictor = LLMPredictor(
    llm=ChatOpenAI(
        temperature=0, model_name="gpt-3.5-turbo", max_tokens=768, client=None
    )
)
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)

qa_prompt_teamplate = """`事前情報`を参考に、後述の`条件`に基づいて、質問に回答してください。`事前情報`でわからない`質問`の場合は、不明であると答えてください。

## 事前情報

{context_str}

## 条件

回答をする際には以下の条件を満たしてください。

* 事前情報で回答できない質問の場合は、回答できないと表示してください。
* 回答が長い場合は要約してください。
* 回答は日本語としてください。

## 質問

[title]に関する質問です。{query_str}
"""
qa_prompt_teamplate = qa_prompt_teamplate.replace("[title]", "ParallelCluster")
qa_prompt = QuestionAnswerPrompt(qa_prompt_teamplate)

refine_prompt_template = """(必要に応じて)`追加の情報`を参照し、`質問`への回答として`現在の回答`が適切か見直して回答してください。

## 質問

[title]に関する質問です。{query_str}

## 現在の回答

{existing_answer}

## 追加の情報

{context_msg}

## 条件

* 回答は日本語としてください。
* 見直しが不要であれば、`現在の回答`の内容だけを出力してください。
* `追加の情報`を参照しても適切ではないと判断した場合は、`不明です`とだけ答えてください。
"""
refine_prompt_template = refine_prompt_template.replace("[title]", "ParallelCluster")
refine_prompt = RefinePrompt(refine_prompt_template)

root_folder = "data"
target_folders = [
    os.path.basename(target_folder) for target_folder in glob.glob(f"{root_folder}/*")
]
target_value_map = {
    target_folder: glob.glob(f"{root_folder}/{target_folder}/*.json")[0]
    for target_folder in target_folders
}


def query(target_folder, query):
    target_file = target_value_map[target_folder]
    index = GPTSimpleVectorIndex.load_from_disk(
        target_file, service_context=service_context
    )
    res = index.query(query, text_qa_template=qa_prompt, refine_template=refine_prompt)

    return res.response  # type: ignore


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            target = gr.Dropdown(target_folders, label="対象")
        with gr.Column(scale=4):
            input = gr.Textbox(label="質問", placeholder="質問を入力してください。")
    with gr.Row():
        output = gr.Markdown()

    input.submit(fn=query, inputs=[target, input], outputs=output)

if __name__ == "__main__":
    demo.launch()