更新DifyQueryRetrieval类的初始化参数，改为使用环境变量获取API密钥和基础URL；优化意图识别示例中的参数传递；调整问题和回答的格式描述；增加请求超时设置。

rag2_0/demo/intent_recognition_example.py

@@ -62,8 +62,8 @@ class QueryRewriteProcessor:
                  api_key: str = None, 
                  base_url: str = None, 
                  model_name: str = None,
-                 dify_api_key: str = "dataset-skLjmPVonjHo119OWNf3kAmY", 
+                 dify_dataset_key: str = None, 
-                 dify_base_url: str = "http://172.20.0.145/v1"):
+                 dify_base_url: str = None):
         """
         初始化查询改写处理器
         
@@ -71,13 +71,17 @@ class QueryRewriteProcessor:
             api_key: API密钥，默认使用环境变量
             base_url: API基础URL，默认使用环境变量
             model_name: 模型名称，默认使用环境变量或默认模型
-            dify_api_key: Dify API密钥
+            dify_dataset_key: Dify API密钥
             dify_base_url: Dify API基础URL
         """
         # 初始化意图识别器
         # 使用asyncio.run()运行异步create方法
         self.recognizer_async = asyncio.run(AsyncIntentRecognizer.create())
-        self.dify_query_retrieval = DifyQueryRetrieval(api_key=dify_api_key, base_url=dify_base_url)
+        if not dify_dataset_key:
+            dify_dataset_key = os.getenv("DIFY_DATASET_KEY")
+        if not dify_base_url:
+            dify_base_url = os.getenv("DIFY_BSAE_URL")
+        self.dify_query_retrieval = DifyQueryRetrieval(dify_dataset_key=dify_dataset_key, dify_base_url=dify_base_url)
     
     def is_retrieved_doc_relevant(self, query: str, retrieved_doc: List[Dict[str, Any]]) -> Dict[str, Any]:
         """
@@ -205,14 +209,13 @@ class QueryRewriteProcessor:
                 classification = result["classification"]
                 original_query = result["rewrite"]["rewrite"]
                 query_list = result["query_expand"]["all"]
-                soft_name = result.get("slot_filling", {}).get("filled_data", {}).get("software_name","")
                 # 将字典转换为Classification对象
                 classification_obj = Classification(**classification)
                 
                 # 根据enable_retrieval参数决定是否进行文档检索
                 retrieved_doc = None
                 if enable_retrieval:
-                    retrieved_doc = self.dify_query_retrieval.retrieve(original_query, query_list, classification_obj, soft_name)
+                    retrieved_doc = self.dify_query_retrieval.retrieve(original_query, query_list, classification_obj, current_softname)
                 
                 # 判断检索文档是否相关
                 relevance_result = {}
@@ -439,9 +442,9 @@ def main():
         for idx, query in enumerate(examples):
             if query.strip() == "":
                 continue
-            query="怎么把一个批次拆分成多个批次工程"
+            query="怎么调整报表顺序"
             conversation_context={
-                "current_softname": "配网计价通D3软件"
+                "current_softname": "储能计价通C1软件"
             }
             # 在调试模式下使用完整的参数
             print(json.dumps(processor.process_query(

rag2_0/dify/DifyQueryRetrieval.py

@@ -27,26 +27,29 @@ class DifyQueryRetrieval:
                                      "西藏造价软件知识(new)","新能源造价知识(new)","配网造价知识(new)","技改造价知识(new)",
                                      "配网造价软件知识(new)"]}
 
-    def __init__(self, api_key: str, base_url: str):
+    def __init__(self, dify_dataset_key: str, dify_base_url: str):
-        self._api_key = api_key
+        self._dify_dataset_key = dify_dataset_key
-        self._base_url = base_url
+        self._dify_base_url = dify_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)
+        client = KnowledgeBaseClient(api_key=self._dify_dataset_key, base_url=self._dify_base_url)
         datasets = client.list_datasets(page_size=50)
         datasets_json = datasets.json()
-        return {dataset["name"]:dataset["id"] for dataset in datasets_json["data"]}
+        return {dataset["name"]:dataset 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])
+            dataset_id = self._datasets_list[dataset_name]["id"]
-            documents = knowledge_base_client.retrieve(query, timeout=300)
+            retrieval_model = self._datasets_list[dataset_name]["retrieval_model_dict"]
+            knowledge_base_client = KnowledgeBaseClient(api_key=self._dify_dataset_key, base_url=self._dify_base_url, dataset_id=dataset_id)
+            
+            documents = knowledge_base_client.retrieve(query, retrieval_model=retrieval_model, 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_id"] = dataset_id
                 retrieved_document["dataset_name"] = dataset_name
             
             return retrieved_documents
@@ -78,6 +81,7 @@ class DifyQueryRetrieval:
 
     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)
+        datasets=["电力建设计价通（2018）软件知识(new)", "主网造价知识(new)", "下载安装注册(new)"]
         if len(datasets) == 0:
             return None
         
@@ -103,6 +107,7 @@ class DifyQueryRetrieval:
         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]]:
+        ssss = self.retrieve_by_dataset("怎么调整报表顺序", "电力建设计价通（2018）软件知识(new)")
         all_documents=[]
         # 使用线程池替代无限制创建线程
         # 设置合理的最大线程数，这里使用min(32, len(query_list) * len(datasets))来限制
@@ -112,7 +117,7 @@ class DifyQueryRetrieval:
             futures = {}
             for query in query_list:
                 for dataset in data_set_list:
-                    if dataset not in self._datasets_list:
+                    if dataset not in list(self._datasets_list.keys()):
                         raise ValueError(f"dataset {dataset} not in datasets_list")
                     
                     futures[executor.submit(self.retrieve_by_dataset, query, dataset)] = query
@@ -140,7 +145,7 @@ class DifyQueryRetrieval:
         
         # 对所有检索出来的文档进行重排序
         time_start = time.time()
-        processed_documents = self.data_post_processor(original_query, deduplicated_documents, top_k)
+        processed_documents = self.data_post_processor("怎么调整报表顺序", deduplicated_documents, top_k)
         time_end = time.time()
         logging.info(f"检索后重排序耗时: {time_end - time_start:.2f}秒")
         
@@ -166,7 +171,7 @@ class DifyQueryRetrieval:
         tasks = []
         for query in query_list:
             for dataset in data_set_list:
-                if dataset not in self._datasets_list:
+                if dataset not in list(self._datasets_list.keys()):
                     logging.error(f"dataset {dataset} not in datasets_list")
                     continue
                 

rag2_0/dify/DifyQueryRetrieval_api.py

@@ -64,7 +64,7 @@ dify_query_retrieval = None
 async def startup_event():
     global dify_query_retrieval
     # 初始化DifyQueryRetrieval实例
-    dify_query_retrieval = DifyQueryRetrieval(api_key="dataset-skLjmPVonjHo119OWNf3kAmY", base_url="http://10.1.16.39/v1")
+    dify_query_retrieval = DifyQueryRetrieval(dify_dataset_key=os.getenv("DIFY_DATASET_KEY"), dify_base_url=os.getenv("DIFY_BSAE_URL"))
     logger.info("DifyQueryRetrieval初始化完成")
 
 # 添加健康检查端点

rag2_0/dify/WorkorderToDify.py

@@ -19,7 +19,7 @@ for index, row in pd_data.iterrows():
     if "存在抱怨" in answer:
         answer = answer.split("存在抱怨")[0]
     
-    content = f"问题：{query}\n解决方案：{answer}"
+    content = f"问题：{query}\n回答：{answer}"
     segments_list.append({
         "content": str(content),
         "answer": "",

rag2_0/dify/dify_client/client.py

@@ -8,7 +8,7 @@ class DifyClient:
         self.api_key = api_key
         self.base_url = base_url
 
-    def _send_request(self, method, endpoint, json=None, params=None, stream=False):
+    def _send_request(self, method, endpoint, json=None, params=None, stream=False, timeout=300):
         headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
@@ -16,7 +16,7 @@ class DifyClient:
 
         url = f"{self.base_url}{endpoint}"
         response = requests.request(
-            method, url, json=json, params=params, headers=headers, stream=stream, verify=False
+            method, url, json=json, params=params, headers=headers, stream=stream, verify=False, timeout=timeout
         )
 
         return response

rag2_0/intent_recognition/DataModels.py

@@ -309,7 +309,7 @@ class IntentAndSlotResult(BaseModel):
 class StepBackPrompt(BaseModel):
     """后退提示数据模型"""
     original_query: str = Field(description="原始查询")
-    can_use_back_prompt: bool = Field(description="原始查询是否可以进行后退提示(True/False),如果原始查询没有限定词或其他限定词语，则不能进行后退提示")
+    can_use_back_prompt: bool = Field(description="原始查询是否可以进行后退提示(true/false),如果原始查询没有限定词或其他限定词语，则不能进行后退提示")
     step_back_query: List[str] = Field(description="后退提示生成的抽象查询(多个)")
 
 class FollowUpQuestions(BaseModel):