## 
# 导入分割文本的工具，并把上面给出的解释分成文档块

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 100,
    chunk_overlap  = 20,
)
explanation = '''Autoencoder（自动编码器）是一种非常有趣的数学模型，就像一个魔法盒子，可以帮助我们理解数据是如何在空间中转换和变换的。

这个魔法盒子里有两个部分：编码器（Encoder）和解码器（Decoder）。

首先，让我们来看一下编码器。它是一个把输入数据（比如一张图片、一段视频或者一篇文章）变得更容易看懂的神奇机器。它将输入数据压缩成一个更小的空间，这样原始数据中的信息就会更加强烈地保留下来。

接下来是解码器。它是一个把编码器压缩后的结果恢复成原始数据的神奇机器。它将编码器得到的结果还原成输入数据（也就是我们刚刚压缩过的数据），这样解码器就得到了和原始数据完全一样的输出。

那么，为什么说Autoencoder能够提高数据处理的效率呢？

因为通过有效的数据压缩和恢复，Autoencoder能够减少数据量，从而更快地处理和分析数据。这就好像把一个大箱子变成了一个更小的箱子，虽然里面东西的总量没有变，但是可以更轻松地拿取和移动箱子。

所以，Autoencoder是一个非常有趣的数学模型，它可以帮助我们更好地理解数据在空间中的转换和变换。
'''

texts = text_splitter.create_documents([explanation])

from langchain.embeddings import HuggingFaceEmbeddings
       
model_name = "nghuyong/ernie-3.0-xbase-zh"
#model_name = "nghuyong/ernie-3.0-nano-zh"
#model_name = "shibing624/text2vec-base-chinese"
#model_name = "GanymedeNil/text2vec-large-chinese"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
hf = HuggingFaceEmbeddings(
   model_name=model_name,
   model_kwargs=model_kwargs,
   encode_kwargs=encode_kwargs,
   cache_folder = "/root/autodl-tmp/ChatGLM2-6B/llm_model"
)
query_result = hf.embed_query(texts[0].page_content)
print(query_result)

from langchain.document_loaders import TextLoader

loader = TextLoader("./index.md")
loader.load()

#pip install pypdf
from langchain.document_loaders import PyPDFLoader

loader = PyPDFLoader("example_data/layout-parser-paper.pdf")
pages = loader.load_and_split()
pages[0]

from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
    # Set a really small chunk size, just to show.
    chunk_size = 100,
    chunk_overlap  = 20,
    length_function = len,
)
texts = text_splitter.create_documents([explanation])
print(texts[0])

from langchain.text_splitter import CharacterTextSplitter
from langchain.document_loaders import TextLoader

loader = TextLoader("/root/autodl-tmp/ChatGLM2-6B/read.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

from langchain.embeddings import HuggingFaceEmbeddings
       
model_name = "nghuyong/ernie-3.0-xbase-zh"
#model_name = "nghuyong/ernie-3.0-nano-zh"
#model_name = "shibing624/text2vec-base-chinese"
#model_name = "GanymedeNil/text2vec-large-chinese"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
hf = HuggingFaceEmbeddings(
   model_name=model_name,
   model_kwargs=model_kwargs,
   encode_kwargs=encode_kwargs,
   cache_folder = "/root/autodl-tmp/ChatGLM2-6B/llm_model"
)

from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS

# 导入分割文本的工具，并把上面给出的解释分成文档块

from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = TextLoader("/root/autodl-tmp/ChatGLM2-6B/read.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
db = FAISS.from_documents(docs, hf)

import os
import shutil
from elasticsearch import Elasticsearch
from langchain.vectorstores import ElasticKnnSearch
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from configs.params import ESParams
from embedding import Embeddings
from typing import Dict


def _default_knn_mapping(dims: int) -> Dict:
    """Generates a default index mapping for kNN search."""
    return {
        "properties": {
            "text": {"type": "text"},
            "vector": {
                "type": "dense_vector",
                "dims": dims,
                "index": True,
                "similarity": "cosine",
            },
        }
    }


def load_file(filepath, chunk_size, chunk_overlap):
    loader = TextLoader(filepath, encoding='utf-8')
    documents = loader.load()
    text_splitter = CharacterTextSplitter(separator='\n', chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = text_splitter.split_documents(documents)
    return docs


class ES:
    def __init__(self, embedding_model_path):
        self.es_params = ESParams()
        self.client = Elasticsearch(['{}:{}'.format(self.es_params.url, self.es_params.port)],
                                    basic_auth=(self.es_params.username, self.es_params.passwd),
                                    verify_certs=False)
        self.embedding = Embeddings(embedding_model_path)
        self.es = ElasticKnnSearch(index_name=self.es_params.index_name, embedding=self.embedding,
                                   es_connection=self.client)

    def doc_upload(self, file_obj, chunk_size, chunk_overlap):
        try:
            if not self.client.indices.exists(index=self.es_params.index_name):
                dims = len(self.embedding.embed_query("test"))
                mapping = _default_knn_mapping(dims)
                self.client.indices.create(index=self.es_params.index_name, body={"mappings": mapping})
            filename = os.path.split(file_obj.name)[-1]
            file_path = 'data/' + filename
            shutil.move(file_obj.name, file_path)
            docs = load_file(file_path, chunk_size, chunk_overlap)
            self.es.add_documents(docs)
            return "插入成功"
        except Exception as e:
            return e

    def exact_search(self, query, top_k):
        result = []
        similar_docs = self.es.similarity_search_with_score(query, k=top_k)
        for i in similar_docs:
            result.append({
                'content': i[0].page_content,
                'source': i[0].metadata['source'],
                'score': i[1]
            })
        return result

    def knn_search(self, query, top_k):
        result = []
        query_vector = self.embedding.embed_query(query)
        similar_docs = self.es.knn_search(query=query, query_vector=query_vector, k=top_k)
        hits = [hit for hit in similar_docs["hits"]["hits"]]
        for i in hits:
            result.append({
                'content': i['_source']['text'],
                'source': i['_source']['metadata']['source'],
                'score': i['_score']
            })
        return result

    def hybrid_search(self, query, top_k, knn_boost):
        result = []
        query_vector = self.embedding.embed_query(query)
        similar_docs = self.es.knn_hybrid_search(query=query, query_vector=query_vector, knn_boost=knn_boost,
                                                 query_boost=1 - knn_boost, k=top_k)
        hits = [hit for hit in similar_docs["hits"]["hits"]]
        for i in hits:
            result.append({
                'content': i['_source']['text'],
                'source': i['_source']['metadata']['source'],
                'score': i['_score']
            })
        result = result[:top_k]
        return result


from configs.params import ESParams
from elasticsearch import Elasticsearch

es_params = ESParams()
index_name = es_params.index_name

# %% 初始化ES对象
client = Elasticsearch(['{}:{}'.format(es_params.url, es_params.port)],
                       basic_auth=(es_params.username, es_params.passwd),
                       verify_certs=False)

# %% 连通测试
client.ping()

# %% 检查索引是否存在
index_exists = client.indices.exists(index=index_name)

# %% 新建索引
response = client.indices.create(index=index_name, body=mapping)

# %% 插入数据
response = client.index(index=index_name, id=document_id, document=data)

# %% 更新
rp = client.update(index=index_name, id=document_id, body={"doc": data})

# %% 检查文档是否存在
document_exists = client.exists(index=index_name, id=document_id)

# %% 根据ID删除文档
response = client.delete(index=index_name, id=document_id)

from langchain.retrievers.multi_query import MultiQueryRetriever

llm = ChatGLM(model_path="/root/autodl-tmp/ChatGLM2-6B/llm_model/models--THUDM--chatglm2-6b/snapshots/8eb45c842594b8473f291d0f94e7bbe86ffc67d8")
llm.load_model()

retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=db.as_retriever(), llm=llm
)
query = "chatglm是什么？"
len(retriever_from_llm.get_relevant_documents(query=query))

from typing import List
from langchain import LLMChain
from pydantic import BaseModel, Field
from langchain.prompts import PromptTemplate
from langchain.output_parsers import PydanticOutputParser


# Output parser will split the LLM result into a list of queries
class LineList(BaseModel):
    # "lines" is the key (attribute name) of the parsed output
    lines: List[str] = Field(description="Lines of text")


class LineListOutputParser(PydanticOutputParser):
    def __init__(self) -> None:
        super().__init__(pydantic_object=LineList)

    def parse(self, text: str) -> LineList:
        lines = text.strip().split("\n")
        return LineList(lines=lines)


output_parser = LineListOutputParser()

QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an AI language model assistant. Your task is to generate five 
    different versions of the given user question to retrieve relevant documents from a vector 
    database. By generating multiple perspectives on the user question, your goal is to help
    the user overcome some of the limitations of the distance-based similarity search. 
    Provide these alternative questions seperated by newlines.
    Original question: {question}""",
)
llm = ChatGLM(model_path="/root/autodl-tmp/ChatGLM2-6B/llm_model/models--THUDM--chatglm2-6b/snapshots/8eb45c842594b8473f291d0f94e7bbe86ffc67d8")
llm.load_model()

# Chain
llm_chain = LLMChain(llm=llm, prompt=QUERY_PROMPT, output_parser=output_parser)

# Other inputs
question = "What are the approaches to Task Decomposition?"

query = "chatglm是什么？"
docs = db.similarity_search(query)

from langchain.chains import RetrievalQA
ruff = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff", retriever=db.as_retriever()
)
ruff.run(query)

from langchain.chains import RetrievalQA
ruff = RetrievalQA.from_chain_type(
    llm=llm, chain_type="refine", retriever=db.as_retriever()
)
ruff.run(query)

from langchain.chains import RetrievalQA
ruff = RetrievalQA.from_chain_type(
    llm=llm, chain_type="map_reduce", retriever=db.as_retriever()
)
ruff.run(query)

# 导入LLMChain并定义一个链，用语言模型和提示作为参数。

from langchain.chains import LLMChain
chain = LLMChain(llm=llm, prompt=prompt)

# 只指定输入变量来运行链。
print(chain.run("autoencoder"))
# 定义一个第二个提示

second_prompt = PromptTemplate(
    input_variables=["ml_concept"],
    template="把{ml_concept}的概念描述转换成用500字向我解释，就像我是一个五岁的孩子一样",
)
chain_two = LLMChain(llm=llm, prompt=second_prompt)

# 用上面的两个链定义一个顺序链：第二个链把第一个链的输出作为输入

from langchain.chains import SimpleSequentialChain
overall_chain = SimpleSequentialChain(chains=[chain, chain_two], verbose=True)

# 只指定第一个链的输入变量来运行链。
explanation = overall_chain.run("autoencoder")
print(explanation)

from langchain.prompts import PromptTemplate

prompt_template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

{context}

Question: {question}
Helpful Answer:"""
PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

chain = load_summarize_chain(llm, chain_type="stuff", prompt=PROMPT)
chain.run(docs)

#把基于向量库的检索生成封装成Tool
ruff = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff", retriever=db.as_retriever()
)

tools = [
    Tool(
        name="Ruff QA System",
        func=ruff.run,
        description="useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.",
    ),
]

#挂到Agent上测试效果
agent = initialize_agent(
    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
)

agent.run(
    "What did biden say about ketanji brown jackson in the state of the union address?"
)