# coding=gbk
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import os

# 定义 Embeddings
embeddings = HuggingFaceEmbeddings(model_name="/mnt/sdb/zhaoyuan/rongrunxiang/acge_text_embedding")


# 向量数据库持久化路径
persist_directory = 'data_base/vector_db/chroma'

# 加载数据库
vectordb = Chroma(
    persist_directory=persist_directory,
    embedding_function=embeddings
)

from LLM import ChatGLM_LLM
modelPath = "/mnt/sdb/zhaoyuan/rongrunxiang/glm-4-9b-chat"
llm = ChatGLM_LLM(model_path = modelPath)
llm.predict("你是谁")

from langchain.prompts import PromptTemplate

# 我们所构造的 Prompt 模板
template = """根据用户的问题，从提供的上下文中找到相应的知识，并据此创作一个适合于3-6岁小孩的绘本文本。尽量生动有趣。
{context}
问题: {question}
文本:"""

# 调用 LangChain 的方法来实例化一个 Template 对象，该对象包含了 context 和 question 两个变量，在实际调用时，这两个变量会被检索到的文档片段和用户提问填充
QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context","question"],template=template)

from langchain.chains import RetrievalQA

qa_chain = RetrievalQA.from_chain_type(llm,retriever=vectordb.as_retriever(),return_source_documents=True,chain_type_kwargs={"prompt":QA_CHAIN_PROMPT})

# 检索问答链回答效果
question = "什么是哺乳动物"
result = qa_chain({"query": question})
print("检索问答链回答 question 的结果：")
print(result["result"])

if "source_documents" in result:
    print("\n找到的文档片段：")
    for doc in result["source_documents"]:
        print(doc)  # 这里假设文档片段是直接可打印的，可能需要根据实际情况调整格式


# 仅 LLM 回答效果
result_2 = llm(question)
print("大模型回答 question 的结果：")
print(result_2)