from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_classic.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate

# 初始化一个大模型
model_name = "模型名称"
llm = ChatOpenAI(model=model_name,
                 base_url='模型url',
                 openai_api_key="模型key",
                 temperature=0)

# 加载PDF文件，并且进行切片
# 指定你的 PDF 文件路径
pdf_path = ""

# 1. 加载 PDF
loader = PyPDFLoader(pdf_path)
documents = loader.load()

# 2. 分割文本
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=2000,  # 每个文本块的大小
    chunk_overlap=400  # 块之间的重叠部分，有助于保持上下文
)
document_chunks = text_splitter.split_documents(documents)
print(f"已将文档拆分为 {len(document_chunks)} 个文本块。")

# 初始化嵌入模型，并且存入向量数据库
embeddings = HuggingFaceEmbeddings(
    model_name="{可以使用本地下载好的embedding模型}"
)
# 创建向量数据库
vector_store = FAISS.from_documents(
    documents=document_chunks,
    embedding=embeddings
)
print("向量数据库构建完成。")



# 自定义提示模板以优化答案质量,模版中设置的context和question是langchain的默认链RetrievalQA固定使用的，因此不能随意修改
custom_prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""请根据以下提供的上下文信息来回答问题。如果上下文中没有明确答案，请说明你不知道，不要编造信息。

上下文：
{context}

问题：{question}

请根据上下文提供答案："""
)
# 创建检索式 QA 链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  # 将检索到的文档"堆叠"后传递给 LLM
    retriever=vector_store.as_retriever(search_kwargs={"k": 3}),  # 检索3个最相关的块
    chain_type_kwargs={"prompt": custom_prompt},
    return_source_documents=True,  # 可选：返回参考的源文档
)

# 提出问题
question = "why langchain？"
result = qa_chain.invoke({"query": question})

# 输出答案
print(f"问题: {question}")
print(f"答案: {result['result']}")

# 如果需要，可以查看参考来源
print("\n参考来源：")
for i, doc in enumerate(result["source_documents"]):
    print(f"{i + 1}. 来源片段: {doc.page_content[:200]}...")  # 打印每个来源片段的前200个字符