优化Dify工具逻辑，调整知识提取和重排序流程，增强API调用的重试机制，更新意图识别API以支持更好的错误处理和日志记录，改进多线程检索功能

rag2_0/dify/DifyQueryRetrieval.py

@@ -100,6 +100,50 @@ class DifyQueryRetrieval:
 
         return processed_documents
 
+    def retrieve_api(self, original_query: str, query_list: List[str],data_set_list: List[str])->List[Dict[str, Any]]:
+        all_documents=[]
+        # 使用线程池替代无限制创建线程
+        # 设置合理的最大线程数，这里使用min(32, len(query_list) * len(data_set_list))来限制
+        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.append(executor.submit(self.retrieve_by_dataset, query, dataset))
+            
+            # 等待所有任务完成
+            for future in as_completed(futures):
+                # 处理可能的异常
+                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)
+        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]]) -> List[Dict[str, Any]]:
         reranker_model = XinferenceReRankerModel()
         documents = [document['segment']['content'] for document in all_documents]
@@ -114,7 +158,7 @@ class DifyQueryRetrieval:
                     "dataset_name": document["dataset_name"],
                     "document_id": document['segment']['document_id'],
                     "document_name": document["segment"]["document"]["name"],
-                    "document_data_source_type": document["segment"]["document"]["data_source_type"],
+                    "data_source_type": document["segment"]["document"]["data_source_type"],
                     "segment_id": document["segment"]["id"],
                     "retriever_from": "api",
                     "score": document.get("score", 0),
@@ -122,6 +166,7 @@ class DifyQueryRetrieval:
                     "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"],
@@ -159,6 +204,7 @@ class DifyQueryRetrieval:
 
     
 if __name__ == "__main__":
-    dify_query_retrieval = DifyQueryRetrieval(api_key="dataset-skLjmPVonjHo119OWNf3kAmY", base_url="https://172.20.0.145/v1")
-    datasets = dify_query_retrieval.retrieve("配网工程计价通D3软件如何新建工程？", Classification(vertical_classification="软件问题", sub_classification="软件功能"), "配网工程计价通D3")
+    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("配网工程计价通D3软件如何新建工程？", ["配网工程计价通D3软件如何新建工程？"], ["流式输出缺失"])
     print(datasets)