"""
1. 目前这个版本是把所有的multiquery当成一个问题提交给大模型。后续可以考虑将每一个问题分别送入大模型，然后得到的多个答案，然后在一并汇总。


"""

# # from langchain.document_loaders import UnstructuredFileLoader
# print('start the importing of UnstructuredFileLoader')
# from langchain_community.document_loaders.unstructured import UnstructuredFileLoader
# filepath = "./joeshi_upload.pdf"
# loader = UnstructuredFileLoader(filepath)
# docs = loader.load()
# print('docs now:', docs)
# print('langchain loader:',loader)
# print('end the importing of UnstructuredFileLoader')
from openai import OpenAI
import openai
import os
import rag_source
import langchain_KB
from langchain.vectorstores import FAISS
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain_community.document_loaders.unstructured import UnstructuredFileLoader ## new version. 
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from dotenv import load_dotenv

client = OpenAI()
load_dotenv()
### 设置openai的API key
os.environ["OPENAI_API_KEY"] = os.environ['user_token']
openai.api_key = os.environ['user_token']

# loader = UnstructuredFileLoader(filepath)
# loader = PyPDFLoader(filepath)
# print('langchain loader:',loader)
# docs = loader.load()
# print('docs now:', docs)

## 文本分割
# text_splitter = CharacterTextSplitter(chunk_size=5000, chunk_overlap=200)
# docs = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200).split_documents(docs)

## 创建向量数据库
# embedding_model_name = 'GanymedeNil/text2vec-large-chinese'
# embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name) ## 这里是联网情况下连接huggingface后使用。
# from langchain.embeddings.openai import OpenAIEmbeddings
# embeddings = OpenAIEmbeddings(disallowed_special=()) ## 可能需要更新了。
# print('langchain embeddings:', embeddings)
# vector_store = FAISS.from_documents(docs, embeddings)
# print(vector_store)
# vector_store.save_local('./joeshi/faiss_index') ## 将矢量库保存在本地指定目录。
# vector_store = FAISS.load_local('./joeshi/faiss_index', embeddings)


### 构建一个multi query的方法来进行大模型查询。可以用chatglm的方法来做, 也可以用openai的方法来做。
def generate_sim_query(orignal_question):
    similar_query_prompt = f"""你是一个AI语言模型，你的任务是生成4个不同的用户问题版本，以便从向量数据库中检索相关文档。通过生成用户问题多角度的表述，你的目标是帮助用户克服基于距离的相似性搜索的一些局限性。请按以下要求生成替代问题，并用隔行排列这些：
    1. 使用与原始问题不同的关键词或同义词来重新表述问题。
    2. 尝试从不同的角度提问，例如更具体、更抽象或更概括。
    3. 考虑问题中隐含的前提或假设，并提出针对这些前提或假设的替代问题。
    4. 结合用户可能拥有的背景知识或先前交互的信息，生成更具针对性的问题。
    用户的原始问题是: {orignal_question}"""
    # similar_query_prompt = """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 separated by newlines.
    # Original question: {orignal_question}""" ## English version.
    
    #### 用大模型来生成相似问。
    # response, history = chatglm.model.chat(chatglm.tokenizer, query=similar_query_prompt) ## 从用langchain的自定义方式来做.
    response = client.chat.completions.create(
                    model="gpt-3.5-turbo-16k",
                    messages=[{"role": "user", "content": f'{similar_query_prompt}'}],
                    stream=False,
                )
    
    ## response.choices[0].message.content 才是ChatGPT反馈的返回的str结果。
    similar_questions = response.choices[0].message.content.split("\n") ## 将反馈结果切割成一个个问题。
    
    return similar_questions


# #### 根据相似问查询，返回查询后的结果以及信息来源。
## 在主程序中会进行大模型的查询，这里应该不需要了。
# def multiQuery_prompt(prompt):
#     prompt = generate_sim_query(prompt)
#     docs = vector_store.similarity_search(prompt, k=3) ##NOTE：注意可能新版本这里的k是小写的，之前的版本可能是大写。
#     source = rag_source.rag_source(docs) ## get the K reference source of the RAG answer, in a designed format. 
    
#     return docs, source