From 4fab8874a55111dc99345bdf54f0fb892f5a954c Mon Sep 17 00:00:00 2001
From: ouyangyouzhang <ouyangyouzhang@booway.com.cn>
Date: Mon, 23 Jun 2025 18:08:39 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9EDifyCompareTest.py=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=E4=BB=A5=E5=AE=9E=E7=8E=B0Dify=E6=96=B0=E6=97=A7?=
 =?UTF-8?q?=E6=B5=81=E7=A8=8B=E5=AF=B9=E6=AF=94=E6=B5=8B=E8=AF=95=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=EF=BC=8C=E5=B9=B6=E6=9B=B4=E6=96=B0.gitignore?=
 =?UTF-8?q?=E4=BB=A5=E5=BF=BD=E7=95=A5=E6=96=B0=E6=B5=8B=E8=AF=95=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore                                    |  1 +
 ...est_dify_chatapi.py => DifyCompareTest.py} |  0
 rag2_0/dify/dify_client/client.py             | 37 +++++++++++++++++++
 3 files changed, 38 insertions(+)
 rename rag2_0/dify/{test_dify_chatapi.py => DifyCompareTest.py} (100%)

diff --git a/.gitignore b/.gitignore
index aaf721e..3606cd9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ rag2_0/demo/Test.py
 data/excel/*.xlsx
 rag2_0/demo/ProfessionalTermAnalyzer.py
 data/logs/*
+rag2_0/dify/Test.py
diff --git a/rag2_0/dify/test_dify_chatapi.py b/rag2_0/dify/DifyCompareTest.py
similarity index 100%
rename from rag2_0/dify/test_dify_chatapi.py
rename to rag2_0/dify/DifyCompareTest.py
diff --git a/rag2_0/dify/dify_client/client.py b/rag2_0/dify/dify_client/client.py
index f28cf5e..972c97b 100755
--- a/rag2_0/dify/dify_client/client.py
+++ b/rag2_0/dify/dify_client/client.py
@@ -457,3 +457,40 @@ class KnowledgeBaseClient(DifyClient):
         data = {"segment": segment_data}
         url = f"/datasets/{self._get_dataset_id()}/documents/{document_id}/segments/{segment_id}"
         return self._send_request("POST", url, json=data, **kwargs)
+
+    def retrieve(
+        self,
+        query: str,
+        retrieval_model: dict | None = None,
+        external_retrieval_model: dict | None = None,
+        **kwargs
+    ):
+        """
+        Retrieve from knowledge base.
+
+        :param query: Query keyword for retrieval
+        :param retrieval_model: Retrieval parameters (optional), example:
+            {
+                "search_method": "semantic_search",  # One of: "keyword_search", "semantic_search", "full_text_search", "hybrid_search"
+                "reranking_enable": True,  # Optional, whether to enable reranking
+                "reranking_mode": {  # Optional, required if reranking_enable is True
+                    "reranking_provider_name": "provider_name",
+                    "reranking_model_name": "model_name"
+                },
+                "weights": 0.5,  # Optional, weight for semantic search in hybrid search mode
+                "top_k": 10,  # Optional, number of results to return
+                "score_threshold_enabled": True,  # Optional, whether to enable score threshold
+                "score_threshold": 0.5  # Optional, score threshold value
+            }
+        :param external_retrieval_model: Unused field
+        :return: Response from the API
+        """
+        data = {"query": query}
+        if retrieval_model is not None:
+            data["retrieval_model"] = retrieval_model
+        if external_retrieval_model is not None:
+            data["external_retrieval_model"] = external_retrieval_model
+            
+        url = f"/datasets/{self._get_dataset_id()}/retrieve"
+        return self._send_request("POST", url, json=data, **kwargs)
+