QueryRewrite/rag2_0/dify/DifyQueryRetrieval.py

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, timeout=300)
            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())
        
        if len(deduplicated_documents) == 0:
            return []
        # 对所有检索出来的文档进行重排序
        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())