1 files changed, 58 insertions, 17 deletions

diff --git a/g4f/requests/__init__.py b/g4f/requests/__init__.py
index 83176557..d4ef9cec 100644
--- a/g4f/requests/__init__.py
+++ b/g4f/requests/__init__.py
@@ -1,18 +1,22 @@
 from __future__ import annotations
 
 from urllib.parse import urlparse
+from typing import Union
+from aiohttp import ClientResponse
+from requests import Response as RequestsResponse
 
 try:
-    from curl_cffi.requests import Session
+    from curl_cffi.requests import Session, Response
     from .curl_cffi import StreamResponse, StreamSession
     has_curl_cffi = True
 except ImportError:
-    from typing import Type as Session
+    from typing import Type as Session, Type as Response
     from .aiohttp import StreamResponse, StreamSession
     has_curl_cffi = False
 
-from ..webdriver import WebDriver, WebDriverSession, bypass_cloudflare, get_driver_cookies
-from ..errors import MissingRequirementsError
+from ..webdriver import WebDriver, WebDriverSession
+from ..webdriver import user_config_dir, bypass_cloudflare, get_driver_cookies
+from ..errors import MissingRequirementsError, RateLimitError, ResponseStatusError
 from .defaults import DEFAULT_HEADERS
 
 def get_args_from_browser(
@@ -20,7 +24,8 @@ def get_args_from_browser(
     webdriver: WebDriver = None,
     proxy: str = None,
     timeout: int = 120,
-    do_bypass_cloudflare: bool = True
+    do_bypass_cloudflare: bool = True,
+    virtual_display: bool = False
 ) -> dict:
     """
     Create a Session object using a WebDriver to handle cookies and headers.
@@ -34,21 +39,37 @@ def get_args_from_browser(
     Returns:
         Session: A Session object configured with cookies and headers from the WebDriver.
     """
-    with WebDriverSession(webdriver, "", proxy=proxy, virtual_display=False) as driver:
+    user_data_dir = "" #user_config_dir(f"g4f-{urlparse(url).hostname}")
+    with WebDriverSession(webdriver, user_data_dir, proxy=proxy, virtual_display=virtual_display) as driver:
         if do_bypass_cloudflare:
             bypass_cloudflare(driver, url, timeout)
-        cookies = get_driver_cookies(driver)
         user_agent = driver.execute_script("return navigator.userAgent")
-    parse = urlparse(url)
+        headers = {
+            **DEFAULT_HEADERS,
+            'referer': url,
+            'user-agent': user_agent,
+        }
+        if hasattr(driver, "requests"):
+            for request in driver.requests:
+                if request.url.startswith(url):
+                    for key, value in request.headers.items():
+                        if key in (
+                            "accept-encoding",
+                            "accept-language",
+                            "user-agent",
+                            "sec-ch-ua",
+                            "sec-ch-ua-platform",
+                            "sec-ch-ua-arch",
+                            "sec-ch-ua-full-version",
+                            "sec-ch-ua-platform-version",
+                            "sec-ch-ua-bitness"
+                        ):
+                            headers[key] = value
+                    break
+        cookies = get_driver_cookies(driver)
     return {
         'cookies': cookies,
-        'headers': {
-            **DEFAULT_HEADERS,
-            'Authority': parse.netloc,
-            'Origin': f'{parse.scheme}://{parse.netloc}',
-            'Referer': url,
-            'User-Agent': user_agent,
-        },
+        'headers': headers,
     }
 
 def get_session_from_browser(url: str, webdriver: WebDriver = None, proxy: str = None, timeout: int = 120) -> Session:
@@ -59,5 +80,25 @@ def get_session_from_browser(url: str, webdriver: WebDriver = None, proxy: str =
         **args,
         proxies={"https": proxy, "http": proxy},
         timeout=timeout,
-        impersonate="chrome110"
-    )
\ No newline at end of file
+        impersonate="chrome"
+    )
+
+async def raise_for_status_async(response: Union[StreamResponse, ClientResponse]):
+    if response.status in (429, 402):
+        raise RateLimitError(f"Response {response.status}: Rate limit reached")
+    text = await response.text() if not response.ok else None
+    if response.status == 403 and "<title>Just a moment...</title>" in text:
+        raise ResponseStatusError(f"Response {response.status}: Cloudflare detected")
+    elif not response.ok:
+        raise ResponseStatusError(f"Response {response.status}: {text}")
+
+def raise_for_status(response: Union[StreamResponse, ClientResponse, Response, RequestsResponse]):
+    if isinstance(response, StreamSession) or isinstance(response, ClientResponse):
+        return raise_for_status_async(response)
+
+    if response.status_code in (429, 402):
+        raise RateLimitError(f"Response {response.status_code}: Rate limit reached")
+    elif response.status_code == 403 and "<title>Just a moment...</title>" in response.text:
+        raise ResponseStatusError(f"Response {response.status_code}: Cloudflare detected")
+    elif not response.ok:
+        raise ResponseStatusError(f"Response {response.status_code}: {response.text}")
\ No newline at end of file