brave.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. import logging
  2. from typing import Optional
  3. import requests
  4. from open_webui.apps.retrieval.web.main import SearchResult, get_filtered_results
  5. from open_webui.env import SRC_LOG_LEVELS
  6. log = logging.getLogger(__name__)
  7. log.setLevel(SRC_LOG_LEVELS["RAG"])
  8. def search_brave(
  9. api_key: str, query: str, count: int, filter_list: Optional[list[str]] = None
  10. ) -> list[SearchResult]:
  11. """Search using Brave's Search API and return the results as a list of SearchResult objects.
  12. Args:
  13. api_key (str): A Brave Search API key
  14. query (str): The query to search for
  15. """
  16. url = "https://api.search.brave.com/res/v1/web/search"
  17. headers = {
  18. "Accept": "application/json",
  19. "Accept-Encoding": "gzip",
  20. "X-Subscription-Token": api_key,
  21. }
  22. params = {"q": query, "count": count}
  23. response = requests.get(url, headers=headers, params=params)
  24. response.raise_for_status()
  25. json_response = response.json()
  26. results = json_response.get("web", {}).get("results", [])
  27. if filter_list:
  28. results = get_filtered_results(results, filter_list)
  29. return [
  30. SearchResult(
  31. link=result["url"], title=result.get("title"), snippet=result.get("snippet")
  32. )
  33. for result in results[:count]
  34. ]