from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.documents import Document
import os

from Config.Config import EMBED_MODEL_NAME, EMBED_BASE_URL, EMBED_API_KEY

# 设置环境变量
os.environ["OPENAI_BASE_URL"] = EMBED_BASE_URL
os.environ["OPENAI_API_KEY"] = EMBED_API_KEY

# 模拟长字符串文档内容
long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。

混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。

混凝土的性能取决于其配合比，包括水灰比、砂率等参数。水灰比是影响混凝土强度的关键因素，较小的水灰比通常会产生更高强度的混凝土。

为了改善混凝土的性能，常常会添加各种外加剂，如减水剂、早强剂、缓凝剂等。此外，还可以使用纤维增强、聚合物改性等技术来提高混凝土的韧性和耐久性。

在施工过程中，混凝土需要适当的养护，以确保其强度正常发展。养护措施包括浇水、覆盖保湿、蒸汽养护等。

随着建筑技术的发展，高性能混凝土、自密实混凝土、再生骨料混凝土等新型混凝土不断涌现，为土木工程领域提供了更多的选择。"""

# 创建文档对象
docs = [Document(page_content=long_text, metadata={"source": "simulated_document"})]
print(f"文档数量：{len(docs)} 个")

# 切割文档
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=200, chunk_overlap=100, add_start_index=True
)
all_splits = text_splitter.split_documents(docs)

print(f"切割后的文档块数量：{len(all_splits)}")

# 嵌入模型
embeddings = OpenAIEmbeddings(model=EMBED_MODEL_NAME)

# 向量存储
vector_store = InMemoryVectorStore(embeddings)
ids = vector_store.add_documents(documents=all_splits)

# 向量查询
results = vector_store.similarity_search(
    "混凝土", k=2
)

# 打印所有查询结果
print("查询结果数量：", len(results))
print("查询结果：")
for i, result in enumerate(results):
    print(f"结果 {i+1}:")
    print(result.page_content)
    print("---")