'commit'

dsSchoolBuddy/ElasticSearch/T2_BgeM3.py

@@ -3,7 +3,7 @@ from langchain_core.documents import Document
 from langchain_core.vectorstores import InMemoryVectorStore
 from langchain_openai import OpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from pydantic import SecretStr  # 导入 SecretStr
+from pydantic import SecretStr
 import requests
 import json
 from Config.Config import (
@@ -12,53 +12,66 @@ from Config.Config import (
 )
 
 
-# 模拟长字符串文档内容
-long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。
+def text_to_vector_db(text: str, chunk_size: int = 200, chunk_overlap: int = 100) -> tuple:
+    """
+    将文本存入向量数据库
     
-混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。
+    参数:
+        text: 要入库的文本
+        chunk_size: 文本分割块大小
+        chunk_overlap: 文本块重叠大小
     
-混凝土的性能取决于其配合比，包括水灰比、砂率等参数。水灰比是影响混凝土强度的关键因素，较小的水灰比通常会产生更高强度的混凝土。
+    返回:
+        tuple: (向量存储对象, 文档数量, 分割后的文档块数量)
+    """
+    # 创建文档对象
+    docs = [Document(page_content=text, metadata={"source": "simulated_document"})]
+    doc_count = len(docs)
+    print(f"文档数量：{doc_count} 个")
 
-为了改善混凝土的性能，常常会添加各种外加剂，如减水剂、早强剂、缓凝剂等。此外，还可以使用纤维增强、聚合物改性等技术来提高混凝土的韧性和耐久性。
+    # 切割文档
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True
+    )
+    all_splits = text_splitter.split_documents(docs)
+    split_count = len(all_splits)
+    print(f"切割后的文档块数量：{split_count}")
 
-在施工过程中，混凝土需要适当的养护，以确保其强度正常发展。养护措施包括浇水、覆盖保湿、蒸汽养护等。
-
-随着建筑技术的发展，高性能混凝土、自密实混凝土、再生骨料混凝土等新型混凝土不断涌现，为土木工程领域提供了更多的选择。"""
-
-# 创建文档对象
-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(
+    # 嵌入模型
+    embeddings = OpenAIEmbeddings(
         model=EMBED_MODEL_NAME,
         base_url=EMBED_BASE_URL,
         api_key=SecretStr(EMBED_API_KEY)  # 包装成 SecretStr 类型
-)
+    )
 
-# 向量存储
-vector_store = InMemoryVectorStore(embeddings)
-ids = vector_store.add_documents(documents=all_splits)
+    # 向量存储
+    vector_store = InMemoryVectorStore(embeddings)
+    ids = vector_store.add_documents(documents=all_splits)
 
-# 向量查询 - 获取更多结果用于重排
-query = "混凝土"
-results = vector_store.similarity_search(query, k=4)  # 获取4个结果用于重排
+    return vector_store, doc_count, split_count
 
-print("向量搜索结果数量：", len(results))
 
-# 存储重排后的文档和分数
-reranked_docs_with_scores = []
+def query_vector_db(vector_store: InMemoryVectorStore, query: str, k: int = 4) -> list:
+    """
+    从向量数据库查询文本
     
-# 调用重排模型
-if len(results) > 1:
+    参数:
+        vector_store: 向量存储对象
+        query: 查询字符串
+        k: 要返回的结果数量
+    
+    返回:
+        list: 重排后的结果列表，每个元素是(文档对象, 可信度分数)的元组
+    """
+    # 向量查询 - 获取更多结果用于重排
+    results = vector_store.similarity_search(query, k=k)
+    print(f"向量搜索结果数量：{len(results)}")
+
+    # 存储重排后的文档和分数
+    reranked_docs_with_scores = []
+
+    # 调用重排模型
+    if len(results) > 1:
         # 准备重排请求数据
         rerank_data = {
             "model": RERANK_MODEL,
@@ -89,18 +102,46 @@ if len(results) > 1:
                 print("警告: 无法识别重排API响应格式")
                 reranked_docs_with_scores = [(doc, 0.0) for doc in results]
 
-        print("重排后结果数量：", len(reranked_docs_with_scores))
+            print(f"重排后结果数量：{len(reranked_docs_with_scores)}")
         except Exception as e:
             print(f"重排模型调用失败: {e}")
             print("将使用原始搜索结果")
             reranked_docs_with_scores = [(doc, 0.0) for doc in results]
-else:
+    else:
         # 只有一个结果，无需重排
         reranked_docs_with_scores = [(doc, 1.0) for doc in results]  # 单个结果可信度设为1.0
 
-# 打印所有查询结果及其可信度
-print("最终查询结果：")
-for i, (result, score) in enumerate(reranked_docs_with_scores):
+    return reranked_docs_with_scores
+
+
+def main():
+    # 模拟长字符串文档内容
+    long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。
+
+混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。
+
+混凝土的性能取决于其配合比，包括水灰比、砂率等参数。水灰比是影响混凝土强度的关键因素，较小的水灰比通常会产生更高强度的混凝土。
+
+为了改善混凝土的性能，常常会添加各种外加剂，如减水剂、早强剂、缓凝剂等。此外，还可以使用纤维增强、聚合物改性等技术来提高混凝土的韧性和耐久性。
+
+在施工过程中，混凝土需要适当的养护，以确保其强度正常发展。养护措施包括浇水、覆盖保湿、蒸汽养护等。
+
+随着建筑技术的发展，高性能混凝土、自密实混凝土、再生骨料混凝土等新型混凝土不断涌现，为土木工程领域提供了更多的选择。"""
+
+    # 调用文本入库功能
+    vector_store, doc_count, split_count = text_to_vector_db(long_text)
+
+    # 调用文本查询功能
+    query = "混凝土"
+    reranked_results = query_vector_db(vector_store, query, k=4)
+
+    # 打印所有查询结果及其可信度
+    print("最终查询结果：")
+    for i, (result, score) in enumerate(reranked_results):
         print(f"结果 {i+1} (可信度: {score:.4f}):")
         print(result.page_content)
         print("---")
+
+
+if __name__ == "__main__":
+    main()