diff options
Diffstat (limited to 'g4f')
-rw-r--r-- | g4f/client/AsyncClient.py | 38 | ||||
-rw-r--r-- | g4f/client/__init__.py | 3 |
2 files changed, 40 insertions, 1 deletions
diff --git a/g4f/client/AsyncClient.py b/g4f/client/AsyncClient.py new file mode 100644 index 00000000..fd2cc353 --- /dev/null +++ b/g4f/client/AsyncClient.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from .Client import Client, Chat, Images, Completions +from .Client import async_iter_response, async_iter_append_model_and_provider +from aiohttp import ClientSession +from typing import Union, AsyncIterator + +class AsyncClient(Client): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.chat = AsyncChat(self) + self._images = AsyncImages(self) + + @property + def images(self) -> 'AsyncImages': + return self._images + +class AsyncCompletions(Completions): + async def async_create(self, *args, **kwargs) -> Union['ChatCompletion', AsyncIterator['ChatCompletionChunk']]: + response = await super().async_create(*args, **kwargs) + async for result in response: + return result + +class AsyncChat(Chat): + def __init__(self, client: AsyncClient): + self.completions = AsyncCompletions(client) + +class AsyncImages(Images): + async def async_generate(self, *args, **kwargs) -> 'ImagesResponse': + return await super().async_generate(*args, **kwargs) + + async def _fetch_image(self, url: str) -> bytes: + async with ClientSession() as session: + async with session.get(url) as resp: + if resp.status == 200: + return await resp.read() + else: + raise Exception(f"Failed to fetch image from {url}, status code {resp.status}") diff --git a/g4f/client/__init__.py b/g4f/client/__init__.py index d1e7e298..0d4685cc 100644 --- a/g4f/client/__init__.py +++ b/g4f/client/__init__.py @@ -1,2 +1,3 @@ from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse -from .client import Client, AsyncClient +from .client import Client +from .client import AsyncClient |