raise_for_status.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. from __future__ import annotations
  2. from typing import Union
  3. from aiohttp import ClientResponse
  4. from requests import Response as RequestsResponse
  5. from ..errors import ResponseStatusError, RateLimitError
  6. from . import Response, StreamResponse
  7. class CloudflareError(ResponseStatusError):
  8. ...
  9. def is_cloudflare(text: str) -> bool:
  10. if "Generated by cloudfront" in text:
  11. return True
  12. elif "<title>Attention Required! | Cloudflare</title>" in text or 'id="cf-cloudflare-status"' in text:
  13. return True
  14. return '<div id="cf-please-wait">' in text or "<title>Just a moment...</title>" in text
  15. def is_openai(text: str) -> bool:
  16. return "<p>Unable to load site</p>" in text
  17. async def raise_for_status_async(response: Union[StreamResponse, ClientResponse], message: str = None):
  18. if response.status in (429, 402):
  19. raise RateLimitError(f"Response {response.status}: Rate limit reached")
  20. message = await response.text() if not response.ok and message is None else message
  21. if response.status == 403 and is_cloudflare(message):
  22. raise CloudflareError(f"Response {response.status}: Cloudflare detected")
  23. elif response.status == 403 and is_openai(message):
  24. raise ResponseStatusError(f"Response {response.status}: Bot are detected")
  25. elif not response.ok:
  26. raise ResponseStatusError(f"Response {response.status}: {message}")
  27. def raise_for_status(response: Union[Response, StreamResponse, ClientResponse, RequestsResponse], message: str = None):
  28. if hasattr(response, "status"):
  29. return raise_for_status_async(response, message)
  30. if response.status_code in (429, 402):
  31. raise RateLimitError(f"Response {response.status_code}: Rate limit reached")
  32. elif response.status_code == 403 and is_cloudflare(response.text):
  33. raise CloudflareError(f"Response {response.status_code}: Cloudflare detected")
  34. elif not response.ok:
  35. raise ResponseStatusError(f"Response {response.status_code}: {response.text if message is None else message}")