from dotenv import load_dotenv
load_dotenv()

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
embedings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
text = [
    "篮球是一项伟大的运动。",
    "带我飞往月球是我最喜欢的歌曲之一。",
    "芝加哥公牛队是我最喜欢的球队。",
    "这是一篇关于芝加哥公牛队的文件。",
    "我非常喜欢去看电影。",
    "芝加哥公牛队以20分的优势赢得了比赛。",
    "天文学是我的另一个兴趣，我常常在晚上观察星空。",
    "《艾尔登之环》是过去15年最好的游戏之一。",
    "迈克尔·乔丹是芝加哥公牛队史最好的球员。",
    "我常常阅读科幻小说，享受其中的幻想世界。",
    "我对环境保护非常关注，参加了一些志愿者活动。"
]
retrieval = Chroma.from_texts(text, embedding=embedings)\
    .as_retriever(search_kwargs={"k": 10})
query = "关于芝加哥公牛队你知道什么?"
# 根据相关性返回文本块
docs = retrieval.get_relevant_documents(query)
print(docs)
print("-------------------------------")
# 提取每个 Document 对象的 page_content 属性
page_contents = [doc.page_content for doc in docs]
# 输出结果
for content in page_contents:
    print(content)