From 7d2ad8c4bf44e2c5b180310a462e9fc90d2ad2ec Mon Sep 17 00:00:00 2001 From: Que Nguyen Date: Mon, 17 Jun 2024 14:34:59 +0700 Subject: [PATCH] Set filter_list as optional param in duckduckgo.py --- backend/apps/rag/search/duckduckgo.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/backend/apps/rag/search/duckduckgo.py b/backend/apps/rag/search/duckduckgo.py index 9342e53e4..8bcb18cc9 100644 --- a/backend/apps/rag/search/duckduckgo.py +++ b/backend/apps/rag/search/duckduckgo.py @@ -1,6 +1,6 @@ import logging -from typing import List -from apps.rag.search.main import SearchResult, filter_by_whitelist +from typing import List, Optional +from apps.rag.search.main import SearchResult, get_filtered_results from duckduckgo_search import DDGS from config import SRC_LOG_LEVELS @@ -8,7 +8,7 @@ log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["RAG"]) -def search_duckduckgo(query: str, count: int, whitelist:List[str]) -> list[SearchResult]: +def search_duckduckgo(query: str, count: int, filter_list: Optional[List[str]] = None) -> list[SearchResult]: """ Search using DuckDuckGo's Search API and return the results as a list of SearchResult objects. Args: @@ -41,7 +41,7 @@ def search_duckduckgo(query: str, count: int, whitelist:List[str]) -> list[Searc snippet=result.get("body"), ) ) - # print(results) - filtered_results = filter_by_whitelist(results, whitelist) + if filter_list: + results = get_filtered_results(results, filter_list) # Return the list of search results - return filtered_results + return results