from langchain.memory import VectorStoreRetrieverMemory
from langchain.embeddings import HuggingFaceBgeEmbeddings
from server.utils import embedding_device, get_model_path
from configs import EMBEDDING_MODEL
import faiss
 
from langchain.docstore import InMemoryDocstore
from langchain.vectorstores import FAISS


# embedding_fn.embed_query("My favorite food is pizza")
embedding_size = 1024 # Dimensions of the OpenAIEmbeddings
index = faiss.IndexFlatL2(embedding_size)
model=EMBEDDING_MODEL
if 'zh' in model:
    # for chinese model
    query_instruction = "为这个句子生成表示以用于检索相关文章："
elif 'en' in model:
    # for english model
    query_instruction = "Represent this sentence for searching relevant passages:"
else:
    # maybe ReRanker or else, just use empty string instead
    query_instruction = ""
embedding_fn = HuggingFaceBgeEmbeddings(model_name=get_model_path(model),
                                                          model_kwargs={'device': embedding_device()},
                                                          query_instruction=query_instruction) 
vectorstore = FAISS(embedding_fn, index, InMemoryDocstore({}), {})

# vectorstore.save_local(folder_path="/data/mnt/nj-1/usr/lixiang/Langchain-Chatchat",index_name = "memory")
# vectorstore.load_local(folder_path="/data/mnt/nj-1/usr/lixiang/Langchain-Chatchat",embeddings=embedding_fn,index_name = "memory")

# In actual usage, you would set `k` to be a higher value, but we use k=1 to show that
# the vector lookup still returns the semantically relevant information
retriever = vectorstore.as_retriever(search_kwargs=dict(k=3))
memory = VectorStoreRetrieverMemory(retriever=retriever)
 
# When added to an agent, the memory object can save pertinent information from conversations or used tools
memory.save_context({"input": "My favorite food is pizza"}, {"output": "thats good to know \n\n\n good"})
memory.save_context({"input": "My favorite sport is soccer"}, {"output": "... \n\n\n good"})
memory.save_context({"input": "I don't the Celtics"}, {"output": "ok \n\n\n good"}) # 

# Notice the first result returned is the memory pertaining to tax help, which the language model deems more semantically relevant
# to a 1099 than the other documents, despite them both containing numbers.
print(memory.load_memory_variables({"prompt": "what sport should i watch?"})["history"])

# print(len(vectorstore.index_to_docstore_id))
# print(vectorstore.index_to_docstore_id[1]) # 6e7874bd-b819-4ecf-b5c4-b5f016013d7d
# print(vectorstore.docstore.search("6e7874bd-b819-4ecf-b5c4-b5f016013d7d"))