external.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. import logging
  2. from typing import Optional, List
  3. import requests
  4. from fastapi import Request
  5. from open_webui.env import SRC_LOG_LEVELS
  6. from open_webui.retrieval.web.main import SearchResult, get_filtered_results
  7. from open_webui.utils.headers import include_user_info_headers
  8. log = logging.getLogger(__name__)
  9. log.setLevel(SRC_LOG_LEVELS["RAG"])
  10. def search_external(
  11. request: Request,
  12. external_url: str,
  13. external_api_key: str,
  14. query: str,
  15. count: int,
  16. filter_list: Optional[List[str]] = None,
  17. user=None,
  18. ) -> List[SearchResult]:
  19. try:
  20. headers = {
  21. "User-Agent": "Open WebUI (https://github.com/open-webui/open-webui) RAG Bot",
  22. "Authorization": f"Bearer {external_api_key}",
  23. }
  24. headers = include_user_info_headers(headers, user)
  25. chat_id = getattr(request.state, "chat_id", None)
  26. if chat_id:
  27. headers["X-OpenWebUI-Chat-Id"] = str(chat_id)
  28. response = requests.post(
  29. external_url,
  30. headers=headers,
  31. json={
  32. "query": query,
  33. "count": count,
  34. },
  35. )
  36. response.raise_for_status()
  37. results = response.json()
  38. if filter_list:
  39. results = get_filtered_results(results, filter_list)
  40. results = [
  41. SearchResult(
  42. link=result.get("link"),
  43. title=result.get("title"),
  44. snippet=result.get("snippet"),
  45. )
  46. for result in results[:count]
  47. ]
  48. log.info(f"External search results: {results}")
  49. return results
  50. except Exception as e:
  51. log.error(f"Error in External search: {e}")
  52. return []