diff --git a/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py b/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py index 7d185508..2e828947 100644 --- a/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py +++ b/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,8 +12,111 @@ from Config.Config import ( ) -# 模拟长字符串文档内容 -long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。 +def text_to_vector_db(text: str, chunk_size: int = 200, chunk_overlap: int = 100) -> tuple: + """ + 将文本存入向量数据库 + + 参数: + 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}") + + # 嵌入模型 + 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) + + return vector_store, doc_count, split_count + + +def query_vector_db(vector_store: InMemoryVectorStore, query: str, k: int = 4) -> list: + """ + 从向量数据库查询文本 + + 参数: + 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, + "query": query, + "documents": [doc.page_content for doc in results], + "top_n": len(results) + } + + # 调用SiliconFlow API进行重排 + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {RERANK_BINDING_API_KEY}" + } + + try: + response = requests.post(RERANK_BASE_URL, headers=headers, data=json.dumps(rerank_data)) + response.raise_for_status() # 检查请求是否成功 + rerank_result = response.json() + + # 处理重排结果，提取relevance_score + if "results" in rerank_result: + for item in rerank_result["results"]: + doc_idx = item.get("index") + score = item.get("relevance_score", 0.0) + if 0 <= doc_idx < len(results): + reranked_docs_with_scores.append((results[doc_idx], score)) + else: + print("警告: 无法识别重排API响应格式") + reranked_docs_with_scores = [(doc, 0.0) for doc in results] + + 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: + # 只有一个结果，无需重排 + reranked_docs_with_scores = [(doc, 1.0) for doc in results] # 单个结果可信度设为1.0 + + return reranked_docs_with_scores + + +def main(): + # 模拟长字符串文档内容 + long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。 混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。 @@ -25,82 +128,20 @@ long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂 随着建筑技术的发展，高性能混凝土、自密实混凝土、再生骨料混凝土等新型混凝土不断涌现，为土木工程领域提供了更多的选择。""" -# 创建文档对象 -docs = [Document(page_content=long_text, metadata={"source": "simulated_document"})] -print(f"文档数量：{len(docs)} 个") + # 调用文本入库功能 + vector_store, doc_count, split_count = text_to_vector_db(long_text) -# 切割文档 -text_splitter = RecursiveCharacterTextSplitter( - chunk_size=200, chunk_overlap=100, add_start_index=True -) -all_splits = text_splitter.split_documents(docs) + # 调用文本查询功能 + query = "混凝土" + reranked_results = query_vector_db(vector_store, query, k=4) -print(f"切割后的文档块数量：{len(all_splits)}") + # 打印所有查询结果及其可信度 + print("最终查询结果：") + for i, (result, score) in enumerate(reranked_results): + print(f"结果 {i+1} (可信度: {score:.4f}):") + print(result.page_content) + print("---") -# 嵌入模型 -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) - -# 向量查询 - 获取更多结果用于重排 -query = "混凝土" -results = vector_store.similarity_search(query, k=4) # 获取4个结果用于重排 - -print("向量搜索结果数量：", len(results)) - -# 存储重排后的文档和分数 -reranked_docs_with_scores = [] - -# 调用重排模型 -if len(results) > 1: - # 准备重排请求数据 - rerank_data = { - "model": RERANK_MODEL, - "query": query, - "documents": [doc.page_content for doc in results], - "top_n": len(results) - } - - # 调用SiliconFlow API进行重排 - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {RERANK_BINDING_API_KEY}" - } - - try: - response = requests.post(RERANK_BASE_URL, headers=headers, data=json.dumps(rerank_data)) - response.raise_for_status() # 检查请求是否成功 - rerank_result = response.json() - - # 处理重排结果，提取relevance_score - if "results" in rerank_result: - for item in rerank_result["results"]: - doc_idx = item.get("index") - score = item.get("relevance_score", 0.0) - if 0 <= doc_idx < len(results): - reranked_docs_with_scores.append((results[doc_idx], score)) - else: - print("警告: 无法识别重排API响应格式") - reranked_docs_with_scores = [(doc, 0.0) for doc in results] - - print("重排后结果数量：", 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: - # 只有一个结果，无需重排 - reranked_docs_with_scores = [(doc, 1.0) for doc in results] # 单个结果可信度设为1.0 - -# 打印所有查询结果及其可信度 -print("最终查询结果：") -for i, (result, score) in enumerate(reranked_docs_with_scores): - print(f"结果 {i+1} (可信度: {score:.4f}):") - print(result.page_content) - print("---") +if __name__ == "__main__": + main()