summaryrefslogtreecommitdiffstats
path: root/g4f/requests/aiohttp.py
diff options
context:
space:
mode:
Diffstat (limited to 'g4f/requests/aiohttp.py')
-rw-r--r--g4f/requests/aiohttp.py26
1 files changed, 21 insertions, 5 deletions
diff --git a/g4f/requests/aiohttp.py b/g4f/requests/aiohttp.py
index d9bd6541..6979b20a 100644
--- a/g4f/requests/aiohttp.py
+++ b/g4f/requests/aiohttp.py
@@ -1,16 +1,20 @@
from __future__ import annotations
-from aiohttp import ClientSession, ClientResponse, ClientTimeout
-from typing import AsyncGenerator, Any
+from aiohttp import ClientSession, ClientResponse, ClientTimeout, BaseConnector
+from typing import AsyncIterator, Any, Optional
-from ..providers.helper import get_connector
from .defaults import DEFAULT_HEADERS
+from ..errors import MissingRequirementsError
class StreamResponse(ClientResponse):
- async def iter_lines(self) -> AsyncGenerator[bytes, None]:
+ async def iter_lines(self) -> AsyncIterator[bytes]:
async for line in self.content:
yield line.rstrip(b"\r\n")
+ async def iter_content(self) -> AsyncIterator[bytes]:
+ async for chunk in self.content.iter_any():
+ yield chunk
+
async def json(self) -> Any:
return await super().json(content_type=None)
@@ -27,4 +31,16 @@ class StreamSession(ClientSession):
response_class=StreamResponse,
connector=get_connector(kwargs.get("connector"), proxies.get("https")),
headers=headers
- ) \ No newline at end of file
+ )
+
+def get_connector(connector: BaseConnector = None, proxy: str = None, rdns: bool = False) -> Optional[BaseConnector]:
+ if proxy and not connector:
+ try:
+ from aiohttp_socks import ProxyConnector
+ if proxy.startswith("socks5h://"):
+ proxy = proxy.replace("socks5h://", "socks5://")
+ rdns = True
+ connector = ProxyConnector.from_url(proxy, rdns=rdns)
+ except ImportError:
+ raise MissingRequirementsError('Install "aiohttp_socks" package for proxy support')
+ return connector \ No newline at end of file