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) +