#!/usr/bin/env python
# coding: utf-8

#! pip install langchain openai chromadb tiktoken pypdf panel

# 注释1:首先，代码导入了所需的库和模块。这包括用于处理 PDF 文件、创建文本嵌入、搜索、问答等的模块。

import os
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
from langchain.document_loaders import TextLoader
from langchain.document_loaders import PyPDFLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
import panel as pn
import tempfile


# In[ ]:


pn.extension('texteditor', template="bootstrap", sizing_mode='stretch_width')
pn.state.template.param.update(
    main_max_width="690px",
    header_background="#F08080",
)


# 注释2:然后，代码设置了一些 Panel 库的参数和小部件，用于创建用户界面。这些小部件包括文件输入框、密码输入框、文本编辑器、按钮、滑块等。这些小部件被组织在一个列中，方便用户输入。


file_input = pn.widgets.FileInput(width=300)

openaikey = pn.widgets.PasswordInput(
    value="", placeholder="Enter your OpenAI API Key here...", width=300
)
prompt = pn.widgets.TextEditor(
    value="", placeholder="Enter your questions here...", height=160, toolbar=False
)
run_button = pn.widgets.Button(name="Run!")

select_k = pn.widgets.IntSlider(
    name="Number of relevant chunks", start=1, end=5, step=1, value=2
)
select_chain_type = pn.widgets.RadioButtonGroup(
    name='Chain type',
    options=['stuff', 'map_reduce', "refine", "map_rerank"]
)

widgets = pn.Row(
    pn.Column(prompt, run_button, margin=5),
    pn.Card(
        "Chain type:",
        pn.Column(select_chain_type, select_k),
        title="Advanced settings", margin=10
    ), width=600
)


#注释3:qa 函数定义了处理 PDF 文件、创建向量嵌入、执行搜索和返回答案的主要逻辑。它接受一个 PDF 文件、查询问题、搜索类型和返回的结果数量作为输入。
#然后，使用 OpenAI 的模型为每个文本块创建一个向量嵌入，这些嵌入被用来创建一个向量存储，用于后续的搜索。接着，使用这个向量存储创建一个检索器，然后使用这个检索器和 OpenAI 的模型创建一个问答链来回答问题。
#最后，函数打印出结果并返回。


def qa(file, query, chain_type, k):
    # load document  加载PDF文件
    loader = PyPDFLoader(file)
    documents = loader.load()
    # split the documents into chunks 将PDF文件分割成小块。
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    texts = text_splitter.split_documents(documents)
    # select which embeddings we want to use 使用 OpenAI 的embeddings模型为每个文本块创建一个向量嵌入
    embeddings = OpenAIEmbeddings()
    # create the VectorStore to use as the index 这些嵌入被用来创建一个向量存储VectorStore，用于后续的搜索。
    db = Chroma.from_documents(texts, embeddings)
    # expose this index in a retriever interface 接着，使用这个向量存储创建一个检索器retriever
    retriever = db.as_retriever(
        search_type="similarity", search_kwargs={"k": k})
    # create a chain to answer questions 然后使用这个检索器和 OpenAI 的模型创建一个问答链来回答问题。
    qa = RetrievalQA.from_chain_type(
        llm=OpenAI(), chain_type=chain_type, retriever=retriever, return_source_documents=True)
    result = qa({"query": query})
    # 最后，函数打印出结果并返回。
    print(result['result'])
    return result


# In[6]:


convos = []  # store all panel objects in a list convos是对话列表的意思

#qa_result 函数是用于处理用户界面输入和调用 qa 函数的函数。它首先从环境变量中获取 OpenAI 的 API 密钥，然后保存用户上传的 PDF 文件。如果用户输入了问题，函数将调用 qa 函数，并将结果添加到对话列表中。
def qa_result(_):
    #首先从环境变量中获取 OpenAI 的 API 密钥
    os.environ["OPENAI_API_KEY"] = openaikey.value

    # save pdf file to a temp file 保存用户上传的 PDF文件为temp.pdf
    if file_input.value is not None:
        file_input.save("/.cache/temp.pdf")

        prompt_text = prompt.value
        if prompt_text:
            result = qa(file="/.cache/temp.pdf", query=prompt_text,
                        chain_type=select_chain_type.value, k=select_k.value)
            convos.extend([
                pn.Row(
                    pn.panel("\U0001F60A", width=10),
                    prompt_text,
                    width=600
                ),
                pn.Row(
                    pn.panel("\U0001F916", width=10),
                    pn.Column(
                        result["result"],
                        "Relevant source text:",
                        pn.pane.Markdown('\n--------------------------------------------------------------------\n'.join(
                            doc.page_content for doc in result["source_documents"]))
                    )
                )
            ])
            # return convos
    return pn.Column(*convos, margin=15, width=575, min_height=400)


# In[7]:创建了一个交互式的 Panel 小部件，当用户点击运行按钮时，会调用 qa_result 函数。
qa_interactive = pn.panel(
    pn.bind(qa_result, run_button),
    loading_indicator=True,
)


# In[8]:创建输出框
output = pn.WidgetBox('*Output will show up here:*',
                      qa_interactive, width=630, scroll=True)


# 界面设计
pn.Column(
    pn.panel.Markdown("""
    ##你可以问我关于你上传的PDF文件的任何信息！
    1) 上传一个PDF文件. 2)输入你的OpenAI API key.这将产生费用  3) 输入问题然后点击"Run".
    """),
    pn.Row(file_input, openaikey),
    output,
    widgets

).servable()