소스 검색

feat: add WEB_SEARCH_RESULT_COUNT to control max number of results

Jun Siang Cheah 1 년 전
부모
커밋
fb8069123e

+ 3 - 3
backend/apps/rag/search/brave.py

@@ -3,7 +3,7 @@ import logging
 import requests
 
 from apps.rag.search.main import SearchResult
-from config import SRC_LOG_LEVELS
+from config import SRC_LOG_LEVELS, WEB_SEARCH_RESULT_COUNT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -22,7 +22,7 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]:
         "Accept-Encoding": "gzip",
         "X-Subscription-Token": api_key,
     }
-    params = {"q": query, "count": 5}
+    params = {"q": query, "count": WEB_SEARCH_RESULT_COUNT}
 
     response = requests.get(url, headers=headers, params=params)
     response.raise_for_status()
@@ -33,5 +33,5 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]:
         SearchResult(
             link=result["url"], title=result.get("title"), snippet=result.get("snippet")
         )
-        for result in results[:5]
+        for result in results[:WEB_SEARCH_RESULT_COUNT]
     ]

+ 2 - 2
backend/apps/rag/search/google_pse.py

@@ -4,7 +4,7 @@ import logging
 import requests
 
 from apps.rag.search.main import SearchResult
-from config import SRC_LOG_LEVELS
+from config import SRC_LOG_LEVELS, WEB_SEARCH_RESULT_COUNT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -27,7 +27,7 @@ def search_google_pse(
         "cx": search_engine_id,
         "q": query,
         "key": api_key,
-        "num": 5,
+        "num": WEB_SEARCH_RESULT_COUNT,
     }
 
     response = requests.request("GET", url, headers=headers, params=params)

+ 2 - 2
backend/apps/rag/search/searxng.py

@@ -3,7 +3,7 @@ import logging
 import requests
 
 from apps.rag.search.main import SearchResult
-from config import SRC_LOG_LEVELS
+from config import SRC_LOG_LEVELS, WEB_SEARCH_RESULT_COUNT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -40,5 +40,5 @@ def search_searxng(query_url: str, query: str) -> list[SearchResult]:
         SearchResult(
             link=result["url"], title=result.get("title"), snippet=result.get("content")
         )
-        for result in sorted_results[:5]
+        for result in sorted_results[:WEB_SEARCH_RESULT_COUNT]
     ]

+ 2 - 2
backend/apps/rag/search/serper.py

@@ -4,7 +4,7 @@ import logging
 import requests
 
 from apps.rag.search.main import SearchResult
-from config import SRC_LOG_LEVELS
+from config import SRC_LOG_LEVELS, WEB_SEARCH_RESULT_COUNT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -35,5 +35,5 @@ def search_serper(api_key: str, query: str) -> list[SearchResult]:
             title=result.get("title"),
             snippet=result.get("description"),
         )
-        for result in results[:5]
+        for result in results[:WEB_SEARCH_RESULT_COUNT]
     ]

+ 2 - 2
backend/apps/rag/search/serpstack.py

@@ -4,7 +4,7 @@ import logging
 import requests
 
 from apps.rag.search.main import SearchResult
-from config import SRC_LOG_LEVELS
+from config import SRC_LOG_LEVELS, WEB_SEARCH_RESULT_COUNT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -39,5 +39,5 @@ def search_serpstack(
         SearchResult(
             link=result["url"], title=result.get("title"), snippet=result.get("snippet")
         )
-        for result in results[:5]
+        for result in results[:WEB_SEARCH_RESULT_COUNT]
     ]

+ 1 - 0
backend/config.py

@@ -549,6 +549,7 @@ BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
 SERPSTACK_API_KEY = os.getenv("SERPSTACK_API_KEY", "")
 SERPSTACK_HTTPS = os.getenv("SERPSTACK_HTTPS", "True").lower() == "true"
 SERPER_API_KEY = os.getenv("SERPER_API_KEY", "")
+WEB_SEARCH_RESULT_COUNT = int(os.getenv("WEB_SEARCH_RESULT_COUNT", "10"))
 
 ####################################
 # Transcribe