Unverified Commit bcb84235 authored by Que Nguyen's avatar Que Nguyen Committed by GitHub
Browse files

Set filter_list as optional param in serply.py

parent 6b8290fa
import json
import logging
from typing import List
from typing import List, Optional
import requests
from urllib.parse import urlencode
from apps.rag.search.main import SearchResult, filter_by_whitelist
from apps.rag.search.main import SearchResult, get_filtered_results
from config import SRC_LOG_LEVELS
log = logging.getLogger(__name__)
......@@ -15,11 +15,11 @@ def search_serply(
api_key: str,
query: str,
count: int,
whitelist:List[str],
hl: str = "us",
limit: int = 10,
device_type: str = "desktop",
proxy_location: str = "US",
filter_list: Optional[List[str]] = None,
) -> list[SearchResult]:
"""Search using serper.dev's API and return the results as a list of SearchResult objects.
......@@ -58,12 +58,13 @@ def search_serply(
results = sorted(
json_response.get("results", []), key=lambda x: x.get("realPosition", 0)
)
filtered_results = filter_by_whitelist(results, whitelist)
if filter_list:
results = get_filtered_results(results, filter_list)
return [
SearchResult(
link=result["link"],
title=result.get("title"),
snippet=result.get("description"),
)
for result in filtered_results[:count]
for result in results[:count]
]
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