Browse Source

Set filter_list as optional param in duckduckgo.py

Que Nguyen 11 months ago
parent
commit
7d2ad8c4bf
1 changed files with 6 additions and 6 deletions
  1. 6 6
      backend/apps/rag/search/duckduckgo.py

+ 6 - 6
backend/apps/rag/search/duckduckgo.py

@@ -1,6 +1,6 @@
 import logging
 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 duckduckgo_search import DDGS
 from config import SRC_LOG_LEVELS
 from config import SRC_LOG_LEVELS
 
 
@@ -8,7 +8,7 @@ log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
 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.
     Search using DuckDuckGo's Search API and return the results as a list of SearchResult objects.
     Args:
     Args:
@@ -41,7 +41,7 @@ def search_duckduckgo(query: str, count: int, whitelist:List[str]) -> list[Searc
                 snippet=result.get("body"),
                 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 the list of search results
-    return filtered_results
+    return results