from langchain_community.docstore import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
import faiss;
from langchain_huggingface import HuggingFaceEmbeddings

# 下载好FAISS

# 准备好向量化的对象
model_name = "BAAI/bge-small-zh-v1.5" # 模型名
model_kwargs = {'device': 'cpu'} # 没有显卡就用cpu，有英伟达显卡写cuda
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity

#  第一次运行，会自动下载模型（去huggingface上下载），下载到hf默认的缓存目录。
hf_embedding = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)
# 1、初始化数据库
# 创建索引
index = faiss.IndexFlatL2(len(hf_embedding.embed_query('Hello world!')));
db = FAISS(
    # 指定好初始化的Embedding的算法模型引用
    embedding_function = hf_embedding,
    # 索引长度
    index = index,
    # 内存存储文档 doc嘛，非关系型数据库
    docstore = InMemoryDocstore(),
    # 字典
    index_to_docstore_id = {}
);

# 2、准备数据
document_1 = Document(
    # 具体数据，key要定死是page_content
    page_content = "今天早餐我吃了巧克力薄煎饼和炒蛋。",
    # 额外数据，随便写，key：value格式
    metadata = {"source": "tweet"}
);
document_2 = Document(
    page_content="明天的天气预报是阴天多云，最高气温62华氏度。",
    metadata={"source": "news"},
)

document_3 = Document(
    page_content="正在用LangChain构建一个激动人心的新项目——快来看看吧！",
    metadata={"source": "tweet"},
)

document_4 = Document(
    page_content="劫匪闯入城市银行，盗走了100万美元现金。",
    metadata={"source": "news"},
)

document_5 = Document(
    page_content="哇！那部电影太精彩了，我已经迫不及待想再看一遍。",
    metadata={"source": "tweet"},
)

document_6 = Document(
    page_content="新iPhone值得这个价格吗？阅读这篇评测一探究竟。",
    metadata={"source": "website"},
)

document_7 = Document(
    page_content="当今世界排名前十的足球运动员。",
    metadata={"source": "website"},
)

document_8 = Document(
    page_content="LangGraph是构建有状态智能体应用的最佳框架！",
    metadata={"source": "tweet"},
)

document_9 = Document(
    page_content="由于对经济衰退的担忧，今日股市下跌500点。",
    metadata={"source": "news"},
)

document_10 = Document(
    page_content="我有种不好的预感，我要被删除了 :(",
    metadata={"source": "tweet"},
)

documents = [
    document_1,
    document_2,
    document_3,
    document_4,
    document_5,
    document_6,
    document_7,
    document_8,
    document_9,
    document_10,
]
# 生成个id，1~documents长度的
ids = [ 'id'+str(i+1) for i in range(len(documents))]
# 数据插入到faiss,并且指定ID，保存到内存~
db.add_documents(documents,ids = ids);

# 把数据库写入磁盘
db.save_local('./faiss_db')


# 语义检索，k=返回几条
# results = db.similarity_search('今天的金融投资新闻', k=2)
results = db.similarity_search('有美食的内容吗', k=2)
for res in results:
    print(type(res))
    print(res.id)
    print(f"* {res.page_content} [{res.metadata}]")