from langchain.prompts.prompt import PromptTemplate
from langchain.llms import OpenAI
from langchain.chains import ChatVectorDBChain

_template = """鑑於以下對話和後續問題，將後續問題改寫為獨立問題。
您可以假設這是有關中草藥和疾病與健康相關的問題。

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:"""
CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)

template = """你是一名 AI 助手，負責回答有關中草藥與健康相關的問題。
您將獲得一份長文檔的以下提取部分和一個問題。 提供對話答案。
如果你不知道答案，就說“我不確定。” 不要試圖編造答案。
如果問題不是關於中草藥與疾病健康的知識，請禮貌地告訴他們你只能回答關於中草藥相關的問題。
Question: {question}
=========
{context}
=========
Answer in Markdown:"""
QA_PROMPT = PromptTemplate(template=template, input_variables=["question", "context"])


def get_chain(vectorstore):
    llm = OpenAI(temperature=0,model_name="gpt-4")
    qa_chain = ChatVectorDBChain.from_llm(
        llm,
        vectorstore,
        condense_question_prompt=CONDENSE_QUESTION_PROMPT,
        return_source_documents=True
    )
    return qa_chain