serply.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. import json
  2. import logging
  3. import requests
  4. from urllib.parse import urlencode
  5. from apps.rag.search.main import SearchResult
  6. from config import SRC_LOG_LEVELS
  7. log = logging.getLogger(__name__)
  8. log.setLevel(SRC_LOG_LEVELS["RAG"])
  9. def search_serply(
  10. api_key: str,
  11. query: str,
  12. count: int,
  13. hl: str = "us",
  14. limit: int = 10,
  15. device_type: str = "desktop",
  16. proxy_location: str = "US"
  17. ) -> list[SearchResult]:
  18. """Search using serper.dev's API and return the results as a list of SearchResult objects.
  19. Args:
  20. api_key (str): A serply.io API key
  21. query (str): The query to search for
  22. hl (str): Host Language code to display results in (reference https://developers.google.com/custom-search/docs/xml_results?hl=en#wsInterfaceLanguages)
  23. limit (int): The maximum number of results to return [10-100, defaults to 10]
  24. """
  25. log.info("Searching with Serply")
  26. url = "https://api.serply.io/v1/search/"
  27. query_payload = {
  28. "q": query,
  29. "language": "en",
  30. "num": limit,
  31. "gl": proxy_location.upper(),
  32. "hl": hl.lower()
  33. }
  34. url = f"{url}{urlencode(query_payload)}"
  35. headers = {
  36. "X-API-KEY": api_key,
  37. "X-User-Agent": device_type,
  38. "User-Agent": "open-webui",
  39. "X-Proxy-Location": proxy_location
  40. }
  41. response = requests.request("GET", url, headers=headers)
  42. response.raise_for_status()
  43. json_response = response.json()
  44. log.info(f"results from serply search: {json_response}")
  45. results = sorted(
  46. json_response.get("results", []), key=lambda x: x.get("realPosition", 0)
  47. )
  48. return [
  49. SearchResult(
  50. link=result["link"],
  51. title=result.get("title"),
  52. snippet=result.get("description"),
  53. )
  54. for result in results[:count]
  55. ]