333 lines
16 KiB
Python
333 lines
16 KiB
Python
import sys
|
|
import os
|
|
import json
|
|
from threading import Thread
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import List, Dict, Any, Optional
|
|
import logging
|
|
import time
|
|
import asyncio
|
|
import httpx
|
|
sys.path.append(os.getcwd())
|
|
|
|
from rag2_0.intent_recognition.DataModels import Classification
|
|
from rag2_0.dify.dify_client.client import DifyClient, KnowledgeBaseClient
|
|
from rag2_0.tool.ModelTool import XinferenceReRankerModel
|
|
class DifyQueryRetrieval:
|
|
|
|
software_to_dataset_map = {"配网工程计价通D3":["下载安装注册(new)","配网造价知识(new)","配网造价软件知识(new)"],
|
|
"新型储能电站建设计价通C1":["下载安装注册(new)","储能C1计价通软件知识(new)","新能源造价知识(new)"],
|
|
"西藏电力工程计价通Z1":["下载安装注册(new)","西藏造价知识(new)","西藏造价软件知识(new)"],
|
|
"技改检修工程计价通T1":["下载安装注册(new)","技改检修工程计价通T1软件知识(new)","技改造价知识(new)"],
|
|
"技改检修清单计价通T1":["下载安装注册(new)","技改检修清单计价通T1软件知识(new)","技改造价知识(new)"],
|
|
"电力建设计价通":["下载安装注册(new)","主网造价知识(new)","电力建设计价通（2018）软件知识(new)"],
|
|
"其他":["下载安装注册(new)","技改检修清单计价通T1软件知识(new)",
|
|
"主网造价知识(new)","西藏造价知识(new)","技改检修工程计价通T1软件知识(new)",
|
|
"电力建设计价通（2018）软件知识(new)","储能C1计价通软件知识(new)",
|
|
"西藏造价软件知识(new)","新能源造价知识(new)","配网造价知识(new)","技改造价知识(new)",
|
|
"配网造价软件知识(new)"]}
|
|
|
|
def __init__(self, api_key: str, base_url: str):
|
|
self._api_key = api_key
|
|
self._base_url = base_url
|
|
self._datasets_list = self.get_datasets_list()
|
|
|
|
def get_datasets_list(self) -> Dict[str, str]:
|
|
client = KnowledgeBaseClient(api_key=self._api_key, base_url=self._base_url)
|
|
datasets = client.list_datasets(page_size=50)
|
|
datasets_json = datasets.json()
|
|
return {dataset["name"]:dataset["id"] for dataset in datasets_json["data"]}
|
|
|
|
def retrieve_by_dataset(self, query: str, dataset_name: str) -> List[Dict[str, Any]]:
|
|
try:
|
|
knowledge_base_client = KnowledgeBaseClient(api_key=self._api_key, base_url=self._base_url, dataset_id=self._datasets_list[dataset_name])
|
|
documents = knowledge_base_client.retrieve(query)
|
|
retrieved_documents = documents.json().get("records", [])
|
|
|
|
# 添加数据集信息
|
|
for retrieved_document in retrieved_documents:
|
|
retrieved_document["dataset_id"] = self._datasets_list[dataset_name]
|
|
retrieved_document["dataset_name"] = dataset_name
|
|
|
|
return retrieved_documents
|
|
except Exception as e:
|
|
logging.error(f"检索数据集 {dataset_name} 时出错: {str(e)}", exc_info=True)
|
|
return []
|
|
|
|
async def retrieve_by_dataset_async(self, query: str, dataset_name: str) -> List[Dict[str, Any]]:
|
|
"""
|
|
异步版本的retrieve_by_dataset方法
|
|
|
|
Args:
|
|
query: 查询字符串
|
|
dataset_name: 数据集名称
|
|
|
|
Returns:
|
|
检索到的文档列表
|
|
"""
|
|
try:
|
|
# 使用asyncio.to_thread包装同步方法
|
|
return await asyncio.to_thread(
|
|
self.retrieve_by_dataset,
|
|
query,
|
|
dataset_name
|
|
)
|
|
except Exception as e:
|
|
logging.error(f"异步检索数据集 {dataset_name} 时出错: {str(e)}", exc_info=True)
|
|
return []
|
|
|
|
def retrieve(self, original_query: str, query_list: List[str], classification: Classification, software_name: str) -> Optional[List[Dict[str, Any]]]:
|
|
datasets = self.get_datasets_by_classification(classification, software_name)
|
|
if len(datasets) == 0:
|
|
return None
|
|
|
|
return self.retrieve_api(original_query, query_list, datasets)
|
|
|
|
async def retrieve_async(self, original_query: str, query_list: List[str], classification: Classification, software_name: str) -> Optional[List[Dict[str, Any]]]:
|
|
"""
|
|
异步版本的retrieve方法
|
|
|
|
Args:
|
|
original_query: 原始查询
|
|
query_list: 查询列表
|
|
classification: 分类信息
|
|
software_name: 软件名称
|
|
|
|
Returns:
|
|
检索到的文档列表
|
|
"""
|
|
datasets = self.get_datasets_by_classification(classification, software_name)
|
|
if len(datasets) == 0:
|
|
return None
|
|
|
|
return await self.retrieve_api_async(original_query, query_list, datasets)
|
|
|
|
def retrieve_api(self, original_query: str, query_list: List[str],data_set_list: List[str], top_k: int = 5)->List[Dict[str, Any]]:
|
|
all_documents=[]
|
|
# 使用线程池替代无限制创建线程
|
|
# 设置合理的最大线程数，这里使用min(32, len(query_list) * len(datasets))来限制
|
|
time_start = time.time()
|
|
max_workers = min(os.cpu_count() * 2, len(query_list) * len(data_set_list))
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {}
|
|
for query in query_list:
|
|
for dataset in data_set_list:
|
|
if dataset not in self._datasets_list:
|
|
raise ValueError(f"dataset {dataset} not in datasets_list")
|
|
|
|
futures[executor.submit(self.retrieve_by_dataset, query, dataset)] = query
|
|
|
|
# 等待所有任务完成
|
|
for future in as_completed(futures.keys()):
|
|
# 处理可能的异常
|
|
try:
|
|
retrieved_documents = future.result()
|
|
all_documents.extend(retrieved_documents)
|
|
except Exception as e:
|
|
logging.error(f"检索过程中发生错误: {str(e)}", exc_info=True)
|
|
time_end = time.time()
|
|
|
|
logging.info(f"检索耗时: {time_end - time_start:.2f}秒")
|
|
# 根据segment_id对文档进行去重
|
|
unique_documents = {}
|
|
for document in all_documents:
|
|
segment_id = document['segment']['id']
|
|
if segment_id not in unique_documents:
|
|
unique_documents[segment_id] = document
|
|
|
|
# 将去重后的文档转换为列表
|
|
deduplicated_documents = list(unique_documents.values())
|
|
|
|
# 对所有检索出来的文档进行重排序
|
|
time_start = time.time()
|
|
processed_documents = self.data_post_processor(original_query, deduplicated_documents, top_k)
|
|
time_end = time.time()
|
|
logging.info(f"检索后重排序耗时: {time_end - time_start:.2f}秒")
|
|
|
|
return processed_documents
|
|
|
|
async def retrieve_api_async(self, original_query: str, query_list: List[str], data_set_list: List[str], top_k: int = 5)->List[Dict[str, Any]]:
|
|
"""
|
|
异步版本的retrieve_api方法，使用asyncio代替线程池
|
|
|
|
Args:
|
|
original_query: 原始查询
|
|
query_list: 查询列表
|
|
data_set_list: 数据集列表
|
|
|
|
Returns:
|
|
检索并重排序后的文档列表
|
|
"""
|
|
all_documents = []
|
|
# 记录开始时间
|
|
time_start = time.time()
|
|
|
|
# 创建异步任务列表
|
|
tasks = []
|
|
for query in query_list:
|
|
for dataset in data_set_list:
|
|
if dataset not in self._datasets_list:
|
|
logging.error(f"dataset {dataset} not in datasets_list")
|
|
continue
|
|
|
|
# 创建异步任务
|
|
task = self.retrieve_by_dataset_async(query, dataset)
|
|
tasks.append(task)
|
|
|
|
# 并发执行所有异步任务
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# 处理结果
|
|
for result in results:
|
|
if isinstance(result, Exception):
|
|
logging.error(f"异步检索过程中发生错误: {str(result)}", exc_info=True)
|
|
else:
|
|
all_documents.extend(result)
|
|
|
|
time_end = time.time()
|
|
logging.info(f"异步检索耗时: {time_end - time_start:.2f}秒")
|
|
|
|
# 根据segment_id对文档进行去重
|
|
unique_documents = {}
|
|
for document in all_documents:
|
|
segment_id = document['segment']['id']
|
|
if segment_id not in unique_documents:
|
|
unique_documents[segment_id] = document
|
|
|
|
# 将去重后的文档转换为列表
|
|
deduplicated_documents = list(unique_documents.values())
|
|
|
|
# 对所有检索出来的文档进行重排序
|
|
time_start = time.time()
|
|
processed_documents = await self.data_post_processor_async(original_query, deduplicated_documents, top_k)
|
|
time_end = time.time()
|
|
logging.info(f"异步检索后重排序耗时: {time_end - time_start:.2f}秒")
|
|
|
|
return processed_documents
|
|
|
|
def data_post_processor(self, query: str, all_documents: List[Dict[str, Any]], top_k: int = 5) -> List[Dict[str, Any]]:
|
|
reranker_model = XinferenceReRankerModel()
|
|
documents = [document['segment']['content'] for document in all_documents]
|
|
reranked_documents = reranker_model.rerank(query, documents, top_k=top_k)
|
|
new_all_documents = []
|
|
|
|
def to_dify_document_format(document: dict)->dict:
|
|
return {
|
|
"metadata": {
|
|
"_source": "knowledge",
|
|
"dataset_id": document["dataset_id"],
|
|
"dataset_name": document["dataset_name"],
|
|
"document_id": document['segment']['document_id'],
|
|
"document_name": document["segment"]["document"]["name"],
|
|
"data_source_type": document["segment"]["document"]["data_source_type"],
|
|
"segment_id": document["segment"]["id"],
|
|
"retriever_from": "api",
|
|
"score": document.get("score", 0),
|
|
"segment_hit_count": document.get("segment", {}).get("hit_count", 0),
|
|
"segment_word_count": document.get("segment", {}).get("word_count", 0),
|
|
"segment_position": document.get("segment", {}).get("position", 0),
|
|
"segment_index_node_hash": document.get("segment", {}).get("index_node_hash", ""),
|
|
"doc_metadata": document.get("segment", {}).get("document", {}).get("doc_metadata", None),
|
|
"position": document["segment"].get("position", 0)
|
|
},
|
|
"title": document["segment"]["document"]["name"],
|
|
"content": document["segment"]["content"]
|
|
}
|
|
|
|
for reranked_document in reranked_documents:
|
|
cur_doc_info = all_documents[reranked_document["index"]]
|
|
cur_doc_info["score"] = reranked_document["score"]
|
|
new_all_documents.append(to_dify_document_format(cur_doc_info))
|
|
return new_all_documents
|
|
|
|
async def data_post_processor_async(self, query: str, all_documents: List[Dict[str, Any]], top_k: int = 5) -> List[Dict[str, Any]]:
|
|
"""
|
|
异步版本的data_post_processor方法
|
|
|
|
Args:
|
|
query: 查询字符串
|
|
all_documents: 待处理的文档列表
|
|
|
|
Returns:
|
|
处理后的文档列表
|
|
"""
|
|
reranker_model = XinferenceReRankerModel()
|
|
documents = [document['segment']['content'] for document in all_documents]
|
|
# 使用异步重排序方法
|
|
reranked_documents = await reranker_model.rerank_async(query, documents, top_k=top_k)
|
|
new_all_documents = []
|
|
|
|
def to_dify_document_format(document: dict)->dict:
|
|
return {
|
|
"metadata": {
|
|
"_source": "knowledge",
|
|
"dataset_id": document["dataset_id"],
|
|
"dataset_name": document["dataset_name"],
|
|
"document_id": document['segment']['document_id'],
|
|
"document_name": document["segment"]["document"]["name"],
|
|
"data_source_type": document["segment"]["document"]["data_source_type"],
|
|
"segment_id": document["segment"]["id"],
|
|
"retriever_from": "api",
|
|
"score": document.get("score", 0),
|
|
"segment_hit_count": document.get("segment", {}).get("hit_count", 0),
|
|
"segment_word_count": document.get("segment", {}).get("word_count", 0),
|
|
"segment_position": document.get("segment", {}).get("position", 0),
|
|
"segment_index_node_hash": document.get("segment", {}).get("index_node_hash", ""),
|
|
"doc_metadata": document.get("segment", {}).get("document", {}).get("doc_metadata", None),
|
|
"position": document["segment"].get("position", 0)
|
|
},
|
|
"title": document["segment"]["document"]["name"],
|
|
"content": document["segment"]["content"]
|
|
}
|
|
|
|
for reranked_document in reranked_documents:
|
|
cur_doc_info = all_documents[reranked_document["index"]]
|
|
cur_doc_info["score"] = reranked_document["score"]
|
|
new_all_documents.append(to_dify_document_format(cur_doc_info))
|
|
return new_all_documents
|
|
|
|
def get_datasets_by_classification(self, classification: Classification, software_name: str) -> List[str]:
|
|
if classification.vertical_classification=="软件问题" or classification.vertical_classification=="业务问题":
|
|
software_name_list = self.software_to_dataset_map.keys()
|
|
cur_software_name = ""
|
|
for software_name_info in software_name_list:
|
|
if software_name_info in software_name:
|
|
cur_software_name = software_name_info
|
|
break
|
|
if cur_software_name == "":
|
|
return self.software_to_dataset_map["其他"]
|
|
else:
|
|
return self.software_to_dataset_map[cur_software_name]
|
|
|
|
if classification.vertical_classification == "安装下载注册":
|
|
if classification.sub_classification in ["后缀名咨询", "软件锁类"]:
|
|
return ["下载安装注册(new)"]
|
|
elif classification.sub_classification == "安装下载类":
|
|
return []
|
|
elif classification.sub_classification == "问题排查":
|
|
return self.software_to_dataset_map["其他"]
|
|
|
|
return self.software_to_dataset_map["其他"]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
dify_query_retrieval = DifyQueryRetrieval(api_key="dataset-skLjmPVonjHo119OWNf3kAmY", base_url="http://10.1.16.39/v1")
|
|
# datasets = dify_query_retrieval.retrieve("配网工程计价通D3软件如何新建工程？", Classification(vertical_classification="软件问题", sub_classification="软件功能"), "配网工程计价通D3")
|
|
# datasets = dify_query_retrieval.retrieve_api("电力建设计价通软件如何批量修改设备价格？",
|
|
# ["电力建设计价通软件如何批量修改设备价格？"],
|
|
# ["电力建设计价通（2018）软件知识(new)"])
|
|
# print(json.dumps(datasets, ensure_ascii=False, indent=2))
|
|
|
|
# 测试异步API
|
|
async def test_async_api():
|
|
datasets = await dify_query_retrieval.retrieve_api_async(
|
|
"电力建设计价通软件如何批量修改设备价格？",
|
|
["电力建设计价通软件如何批量修改设备价格？"],
|
|
["电力建设计价通（2018）软件知识(new)"]
|
|
)
|
|
print("异步API测试结果:")
|
|
print(json.dumps(datasets, ensure_ascii=False, indent=2))
|
|
|
|
# 如果需要测试异步API，取消下面的注释
|
|
import asyncio
|
|
asyncio.run(test_async_api()) |