from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader, WebBaseLoader

# TextLoader
# 指定编码
loader = TextLoader("./input/金庸武侠小说全集/射雕英雄传.txt", encoding="utf-8")
loader.load()
# 自动推断
# python -m pip install chardet
loader = TextLoader("./input/金庸武侠小说全集/射雕英雄传.txt", autodetect_encoding=True)
loader.load()

# PyPDFLoader
# python -m pip install pypdf
loader = PyPDFLoader("./input/文学作品/追风筝的人.pdf")
loader.load()

# WebBaseLoader
# python -m pip install beautifulsoup4
loader = WebBaseLoader(web_paths=('https://blog.yuanpei.me',), bs_kwargs={})
loader.load()

from langchain_community.document_loaders import DirectoryLoader

loader = DirectoryLoader("./posts/", glob="*.md", loader_kwargs={}, show_progress=True, silent_errors=True)

from langchain_community.document_loaders import DirectoryLoader, TextLoader

loader = DirectoryLoader("./posts/", glob="*.md", loader_cls=TextLoader, loader_kwargs={}, show_progress=True, silent_errors=True)

from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter

text_splitter = CharacterTextSplitter(separator = "\n\n", chunk_size = 600, chunk_overlap = 100, length_function = len)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

documents = loader.load()。
text_splitter.split_documents(documents)

documents = loader.load_and_split(text_splitter)

import os
from langchain_community.embeddings import HuggingFaceEmbeddings

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
embeddings = HuggingFaceEmbeddings(model_name="GanymedeNil/text2vec-large-chinese")

# 单个句子
print(embeddings.embed_query("人生若只如初见，何事秋风悲画扇"))
# 多个句子
print(embeddings.embed_documents(["人生若只如初见，何事秋风悲画扇","等闲变却故人心，却道故人心易变"]))

from langchain_community.vectorstores.faiss import FAISS
import pickle

# python -m pip install faiss-cpu
vector_store = FAISS.from_documents(documents, embeddings)
with open('./output/<Your-Persistence-Path>', "wb") as f:
    pickle.dump(vector_store, f)

with open("./output/个人博客.pkl", "rb") as f:
    vectorstore = pickle.load(f)

    query = 'Envoy 在微服务中的应用场景有哪些'
    topK = 3

    # 普通的相似性检索
    result = vectorstore.similarity_search(query=query, k=topK)
    # 带有分数的相似性检索
    result = vectorstore.similarity_search_with_relevance_scores(query=query, k=topK)
    # 最大边界相关性检索
    result = vectorstore.max_marginal_relevance_search(query=query, k=topK)

from langchain.prompts.prompt import PromptTemplate

# 定义提示词模板
PROMPT_TEMPLATE = """
You are a helpful AI bot. Your name is {name}.
Please answer the question only based on the following context:

{context}

If the question is about your identity or role or name, answer '{name}' directly, no need to refer to the context.
If the context is not enough to support the generation of an answer, Please return "I'm sorry, I can't answer your question." immediately.
You have an opportunity to refine the existing answer (only if needed) with current context.
You must always answer the question in Chinese.

Queation: {question}
"""

# 根据问题做相似性检索
question = "Envoy在微服务中都有哪些应用场景"
documents = vector_store.similarity_search(question, k=3)
context = '\n\n'.join([document.page_content for document in documents])

# 从模板创建提示词并填充模板
prompt = PromptTemplate.from_template(PROMPT_TEMPLATE)
query = prompt.format(question=question, name="ChatGPT", context=context)

from langchain_core.prompts import ChatPromptTemplate

# 定义提示词模板
PROMPT_TEMPLATE = """
You are a helpful AI bot. Your name is {name}.
Please answer the question only based on the following context:

{context}

If the question is about your identity or role or name, answer '{name}' directly, no need to refer to the context.
If the context is not enough to support the generation of an answer, Please return "I'm sorry, I can't answer your question." immediately.
You have an opportunity to refine the existing answer (only if needed) with current context.
You must always answer the question in Chinese. 
"""

# 根据问题做相似性检索
question = "Envoy在微服务中都有哪些应用场景"
documents = vector_store.similarity_search(question, k=3)
context = '\n\n'.join([document.page_content for document in documents])

# 从模板创建提示词并填充模板
prompt = ChatPromptTemplate.from_messages([
    ("system", PROMPT_TEMPLATE),
    ("user", "{question}"),
])
query = prompt.format(question=question, name="ChatGPT", context=context)

from langchain.chains import LLMChain
from langchain_community.llms import OpenAI
from langchain_core.prompts import PromptTemplate

prompt_template = "Tell me a {adjective} joke"
prompt = PromptTemplate(
    input_variables=["adjective"], 
    template=prompt_template
)
llm = LLMChain(llm=OpenAI(), prompt=prompt)

# 基本的 Q&A
def get_basic_qa_chain(baseUrl='', apiKey='', storeFilePath=''):
    llm = ChatOpenAI(
        model_name="gpt-3.5-turbo", 
        temperature=0.75, 
        openai_api_base=baseUrl, 
        openai_api_key=apiKey,
        streaming=True
    )
    retriever = load_retriever(storeFilePath)
    chain = RetrievalQA.from_chain_type(
        llm=llm, 
        chain_type="stuff",
        retriever=retriever, 
        return_source_documents=True,
    )
    return chain

# 对话式检索
def get_conversational_retrieval_chain(baseUrl='', apiKey='', storeFilePath=''):
    llm = ChatOpenAI(
        model_name="gpt-3.5-turbo", 
        temperature=0.75, 
        openai_api_base=baseUrl, 
        openai_api_key=apiKey,
        streaming=True
    )
    retriever = load_retriever(storeFilePath)
    memory = ConversationBufferMemory(
        memory_key="chat_history", 
        return_messages=True, 
        input_key="question", 
        output_key="source_documents"
    )
    chain = ConversationalRetrievalChain.from_llm(
        llm=llm, 
        retriever=retriever, 
        memory=memory, 
        return_source_documents=True, 
    )
    return chain

# 构造 Chain，使用本地模型
chain = get_conversational_retrieval_chain_("http://localhost:8080/v1/", 'sk-1234567890',  './output/个人博客.pkl')

# 根据问题做相似性检索
question = "Envoy在微服务中都有哪些应用场景"
documents = vector_store.similarity_search(question, k=3)
context = '\n\n'.join([document.page_content for document in documents])

# 从模板创建提示词并填充模板
prompt = ChatPromptTemplate.from_messages([
    ("system", PROMPT_TEMPLATE),
    ("user", "{question}"),
])
query = prompt.format(question=question, name="ChatGPT", context=context)

# 方式一： 同步调用
result = chain.invoke(query)

# 方式二： 流式调用
result = chain.stream(query)