jina_search.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. import logging
  2. import requests
  3. from open_webui.retrieval.web.main import SearchResult
  4. from open_webui.env import SRC_LOG_LEVELS
  5. from yarl import URL
  6. log = logging.getLogger(__name__)
  7. log.setLevel(SRC_LOG_LEVELS["RAG"])
  8. def search_jina(api_key: str, query: str, count: int) -> list[SearchResult]:
  9. """
  10. Search using Jina's Search API and return the results as a list of SearchResult objects.
  11. Args:
  12. query (str): The query to search for
  13. count (int): The number of results to return
  14. Returns:
  15. list[SearchResult]: A list of search results
  16. """
  17. jina_search_endpoint = "https://s.jina.ai/"
  18. headers = {
  19. "Accept": "application/json",
  20. "Content-Type": "application/json",
  21. "Authorization": api_key,
  22. "X-Retain-Images": "none"
  23. }
  24. payload = {
  25. "q": query,
  26. "count": count if count <= 10 else 10
  27. }
  28. url = str(URL(jina_search_endpoint))
  29. response = requests.post(url, headers=headers, json=payload)
  30. response.raise_for_status()
  31. data = response.json()
  32. results = []
  33. for result in data["data"]:
  34. results.append(
  35. SearchResult(
  36. link=result["url"],
  37. title=result.get("title"),
  38. snippet=result.get("content"),
  39. )
  40. )
  41. return results