Commit fb806912 authored by Jun Siang Cheah's avatar Jun Siang Cheah
Browse files

feat: add WEB_SEARCH_RESULT_COUNT to control max number of results

parent 635951b5
...@@ -3,7 +3,7 @@ import logging ...@@ -3,7 +3,7 @@ import logging
import requests import requests
from apps.rag.search.main import SearchResult 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 = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"]) log.setLevel(SRC_LOG_LEVELS["RAG"])
...@@ -22,7 +22,7 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]: ...@@ -22,7 +22,7 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]:
"Accept-Encoding": "gzip", "Accept-Encoding": "gzip",
"X-Subscription-Token": api_key, "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 = requests.get(url, headers=headers, params=params)
response.raise_for_status() response.raise_for_status()
...@@ -33,5 +33,5 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]: ...@@ -33,5 +33,5 @@ def search_brave(api_key: str, query: str) -> list[SearchResult]:
SearchResult( SearchResult(
link=result["url"], title=result.get("title"), snippet=result.get("snippet") link=result["url"], title=result.get("title"), snippet=result.get("snippet")
) )
for result in results[:5] for result in results[:WEB_SEARCH_RESULT_COUNT]
] ]
...@@ -4,7 +4,7 @@ import logging ...@@ -4,7 +4,7 @@ import logging
import requests import requests
from apps.rag.search.main import SearchResult 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 = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"]) log.setLevel(SRC_LOG_LEVELS["RAG"])
...@@ -27,7 +27,7 @@ def search_google_pse( ...@@ -27,7 +27,7 @@ def search_google_pse(
"cx": search_engine_id, "cx": search_engine_id,
"q": query, "q": query,
"key": api_key, "key": api_key,
"num": 5, "num": WEB_SEARCH_RESULT_COUNT,
} }
response = requests.request("GET", url, headers=headers, params=params) response = requests.request("GET", url, headers=headers, params=params)
......
...@@ -3,7 +3,7 @@ import logging ...@@ -3,7 +3,7 @@ import logging
import requests import requests
from apps.rag.search.main import SearchResult 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 = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"]) log.setLevel(SRC_LOG_LEVELS["RAG"])
...@@ -40,5 +40,5 @@ def search_searxng(query_url: str, query: str) -> list[SearchResult]: ...@@ -40,5 +40,5 @@ def search_searxng(query_url: str, query: str) -> list[SearchResult]:
SearchResult( SearchResult(
link=result["url"], title=result.get("title"), snippet=result.get("content") 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]
] ]
...@@ -4,7 +4,7 @@ import logging ...@@ -4,7 +4,7 @@ import logging
import requests import requests
from apps.rag.search.main import SearchResult 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 = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"]) log.setLevel(SRC_LOG_LEVELS["RAG"])
...@@ -35,5 +35,5 @@ def search_serper(api_key: str, query: str) -> list[SearchResult]: ...@@ -35,5 +35,5 @@ def search_serper(api_key: str, query: str) -> list[SearchResult]:
title=result.get("title"), title=result.get("title"),
snippet=result.get("description"), snippet=result.get("description"),
) )
for result in results[:5] for result in results[:WEB_SEARCH_RESULT_COUNT]
] ]
...@@ -4,7 +4,7 @@ import logging ...@@ -4,7 +4,7 @@ import logging
import requests import requests
from apps.rag.search.main import SearchResult 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 = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"]) log.setLevel(SRC_LOG_LEVELS["RAG"])
...@@ -39,5 +39,5 @@ def search_serpstack( ...@@ -39,5 +39,5 @@ def search_serpstack(
SearchResult( SearchResult(
link=result["url"], title=result.get("title"), snippet=result.get("snippet") link=result["url"], title=result.get("title"), snippet=result.get("snippet")
) )
for result in results[:5] for result in results[:WEB_SEARCH_RESULT_COUNT]
] ]
...@@ -549,6 +549,7 @@ BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "") ...@@ -549,6 +549,7 @@ BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
SERPSTACK_API_KEY = os.getenv("SERPSTACK_API_KEY", "") SERPSTACK_API_KEY = os.getenv("SERPSTACK_API_KEY", "")
SERPSTACK_HTTPS = os.getenv("SERPSTACK_HTTPS", "True").lower() == "true" SERPSTACK_HTTPS = os.getenv("SERPSTACK_HTTPS", "True").lower() == "true"
SERPER_API_KEY = os.getenv("SERPER_API_KEY", "") SERPER_API_KEY = os.getenv("SERPER_API_KEY", "")
WEB_SEARCH_RESULT_COUNT = int(os.getenv("WEB_SEARCH_RESULT_COUNT", "10"))
#################################### ####################################
# Transcribe # Transcribe
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment