1 files changed, 162 insertions, 160 deletions

diff --git a/g4f/client/async_client.py b/g4f/client/async_client.py
index 3ac77b41..2fe4640b 100644
--- a/g4f/client/async_client.py
+++ b/g4f/client/async_client.py
@@ -1,27 +1,32 @@
 from __future__ import annotations
 
-import os
 import time
 import random
 import string
-import logging
 import asyncio
-from typing import Union, AsyncIterator
-from ..providers.base_provider import AsyncGeneratorProvider
-from ..image import ImageResponse, to_image, to_data_uri
-from ..typing import Messages, ImageType
-from ..providers.types import BaseProvider, ProviderType, FinishReason
-from ..providers.conversation import BaseConversation
-from ..image import ImageResponse as ImageProviderResponse
-from ..errors import NoImageResponseError
-from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse
-from .image_models import ImageModels
-from .types import IterResponse, ImageProvider
+import base64
+from aiohttp import ClientSession, BaseConnector
+
 from .types import Client as BaseClient
-from .service import get_model_and_provider, get_last_provider
-from .helper import find_stop, filter_json, filter_none
-from ..models import ModelUtils
-from ..Provider import IterListProvider
+from .types import ProviderType, FinishReason
+from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse, Image
+from .types import AsyncIterResponse, ImageProvider
+from .image_models import ImageModels
+from .helper import filter_json, find_stop, filter_none, cast_iter_async
+from .service import get_last_provider, get_model_and_provider
+from ..Provider import ProviderUtils
+from ..typing import Union, Messages, AsyncIterator, ImageType
+from ..errors import NoImageResponseError, ProviderNotFoundError
+from ..requests.aiohttp import get_connector
+from ..providers.conversation import BaseConversation
+from ..image import ImageResponse as ImageProviderResponse, ImageDataResponse
+
+try:
+    anext
+except NameError:
+    async def anext(iter):
+        async for chunk in iter:
+            return chunk
 
 async def iter_response(
     response: AsyncIterator[str],
@@ -29,37 +34,30 @@ async def iter_response(
     response_format: dict = None,
     max_tokens: int = None,
     stop: list = None
-) -> AsyncIterator[ChatCompletion | ChatCompletionChunk]:
+) -> AsyncIterResponse:
     content = ""
     finish_reason = None
     completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
-
-    async for idx, chunk in enumerate(response):
+    count: int = 0
+    async for chunk in response:
         if isinstance(chunk, FinishReason):
             finish_reason = chunk.reason
             break
         elif isinstance(chunk, BaseConversation):
             yield chunk
             continue
-        
         content += str(chunk)
-        
-        if max_tokens is not None and idx + 1 >= max_tokens:
+        count += 1
+        if max_tokens is not None and count >= max_tokens:
             finish_reason = "length"
-        
-        first, content, chunk = find_stop(stop, content, chunk if stream else None)
-        
+        first, content, chunk = find_stop(stop, content, chunk)
         if first != -1:
             finish_reason = "stop"
-        
         if stream:
             yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
-        
         if finish_reason is not None:
             break
-    
     finish_reason = "stop" if finish_reason is None else finish_reason
-    
     if stream:
         yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
     else:
@@ -68,12 +66,12 @@ async def iter_response(
                 content = filter_json(content)
         yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
 
-async def iter_append_model_and_provider(response: AsyncIterator) -> AsyncIterator:
+async def iter_append_model_and_provider(response: AsyncIterResponse) -> AsyncIterResponse:
     last_provider = None
     async for chunk in response:
         last_provider = get_last_provider(True) if last_provider is None else last_provider
         chunk.model = last_provider.get("model")
-        chunk.provider = last_provider.get("name")
+        chunk.provider =  last_provider.get("name")
         yield chunk
 
 class AsyncClient(BaseClient):
@@ -82,32 +80,59 @@ class AsyncClient(BaseClient):
         provider: ProviderType = None,
         image_provider: ImageProvider = None,
         **kwargs
-    ) -> None:
+    ):
         super().__init__(**kwargs)
         self.chat: Chat = Chat(self, provider)
-        self._images: Images = Images(self, image_provider)
+        self.images: Images = Images(self, image_provider)
 
-    @property
-    def images(self) -> Images:
-        return self._images
+def create_response(
+    messages: Messages,
+    model: str,
+    provider: ProviderType = None,
+    stream: bool = False,
+    proxy: str = None,
+    max_tokens: int = None,
+    stop: list[str] = None,
+    api_key: str = None,
+    **kwargs
+):
+    has_asnyc = hasattr(provider, "create_async_generator")
+    if has_asnyc:
+        create = provider.create_async_generator
+    else:
+        create = provider.create_completion
+    response = create(
+        model, messages,
+        stream=stream,            
+        **filter_none(
+            proxy=proxy,
+            max_tokens=max_tokens,
+            stop=stop,
+            api_key=api_key
+        ),
+        **kwargs
+    )
+    if not has_asnyc:
+        response = cast_iter_async(response)
+    return response
 
-class Completions:
+class Completions():
     def __init__(self, client: AsyncClient, provider: ProviderType = None):
         self.client: AsyncClient = client
         self.provider: ProviderType = provider
 
-    async def create(
+    def create(
         self,
         messages: Messages,
         model: str,
         provider: ProviderType = None,
         stream: bool = False,
         proxy: str = None,
-        response_format: dict = None,
         max_tokens: int = None,
         stop: Union[list[str], str] = None,
         api_key: str = None,
-        ignored: list[str] = None,
+        response_format: dict = None,
+        ignored  : list[str] = None,
         ignore_working: bool = False,
         ignore_stream: bool = False,
         **kwargs
@@ -118,156 +143,133 @@ class Completions:
             stream,
             ignored,
             ignore_working,
-            ignore_stream,
+            ignore_stream
         )
-
         stop = [stop] if isinstance(stop, str) else stop
-        
-        response = await provider.create_completion(
-            model,
-            messages,
-            stream=stream,
-            **filter_none(
-                proxy=self.client.get_proxy() if proxy is None else proxy,
-                max_tokens=max_tokens,
-                stop=stop,
-                api_key=self.client.api_key if api_key is None else api_key
-            ),
+        response = create_response(
+            messages, model,
+            provider, stream,
+            proxy=self.client.get_proxy() if proxy is None else proxy,
+            max_tokens=max_tokens,
+            stop=stop,
+            api_key=self.client.api_key if api_key is None else api_key,
             **kwargs
         )
-        
         response = iter_response(response, stream, response_format, max_tokens, stop)
         response = iter_append_model_and_provider(response)
-        
-        return response if stream else await anext(response)
+        return response if stream else anext(response)
 
-class Chat:
+class Chat():
     completions: Completions
 
     def __init__(self, client: AsyncClient, provider: ProviderType = None):
         self.completions = Completions(client, provider)
 
-async def iter_image_response(response: AsyncIterator) -> Union[ImagesResponse, None]:
-    logging.info("Starting iter_image_response")
+async def iter_image_response(
+    response: AsyncIterator,
+    response_format: str = None,
+    connector: BaseConnector = None,
+    proxy: str = None
+) -> Union[ImagesResponse, None]:
     async for chunk in response:
-        logging.info(f"Processing chunk: {chunk}")
         if isinstance(chunk, ImageProviderResponse):
-            logging.info("Found ImageProviderResponse")
-            return ImagesResponse([Image(image) for image in chunk.get_list()])
-    
-    logging.warning("No ImageProviderResponse found in the response")
-    return None
-
-async def create_image(client: AsyncClient, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:
-    logging.info(f"Creating image with provider: {provider}, model: {model}, prompt: {prompt}")
+            if response_format == "b64_json":
+                async with ClientSession(
+                    connector=get_connector(connector, proxy),
+                    cookies=chunk.options.get("cookies")
+                ) as session:
+                    async def fetch_image(image):
+                        async with session.get(image) as response:
+                            return base64.b64encode(await response.content.read()).decode()
+                    images = await asyncio.gather(*[fetch_image(image) for image in chunk.get_list()])
+                return ImagesResponse([Image(None, image, chunk.alt) for image in images], int(time.time()))
+            return ImagesResponse([Image(image, None, chunk.alt) for image in chunk.get_list()], int(time.time()))
+        elif isinstance(chunk, ImageDataResponse):
+            return ImagesResponse([Image(None, image, chunk.alt) for image in chunk.get_list()], int(time.time()))
 
+def create_image(provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:
     if isinstance(provider, type) and provider.__name__ == "You":
         kwargs["chat_mode"] = "create"
     else:
-        prompt = f"create an image with: {prompt}"
-    
-    response = await provider.create_completion(
+        prompt = f"create a image with: {prompt}"
+    return provider.create_async_generator(
         model,
         [{"role": "user", "content": prompt}],
         stream=True,
-        proxy=client.get_proxy(),
         **kwargs
     )
-    
-    logging.info(f"Response from create_completion: {response}")
-    return response
 
-class Images:
+class Images():
     def __init__(self, client: AsyncClient, provider: ImageProvider = None):
         self.client: AsyncClient = client
         self.provider: ImageProvider = provider
         self.models: ImageModels = ImageModels(client)
 
-    async def generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse:
-        logging.info(f"Starting asynchronous image generation for model: {model}, prompt: {prompt}")
-        provider = self.models.get(model, self.provider)
-        if provider is None:
-            raise ValueError(f"Unknown model: {model}")
-
-        logging.info(f"Provider: {provider}")
-
-        if isinstance(provider, IterListProvider):
-            if provider.providers:
-                provider = provider.providers[0]
-                logging.info(f"Using first provider from IterListProvider: {provider}")
+    def get_provider(self, model: str, provider: ProviderType = None):
+        if isinstance(provider, str):
+            if provider in ProviderUtils.convert:
+                provider = ProviderUtils.convert[provider]
             else:
-                raise ValueError(f"IterListProvider for model {model} has no providers")
-
-        if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
-            logging.info("Using AsyncGeneratorProvider")
-            messages = [{"role": "user", "content": prompt}]
-            async for response in provider.create_async_generator(model, messages, **kwargs):
-                if isinstance(response, ImageResponse):
-                    return self._process_image_response(response)
-                elif isinstance(response, str):
-                    image_response = ImageResponse([response], prompt)
-                    return self._process_image_response(image_response)
-        elif hasattr(provider, 'create'):
-            logging.info("Using provider's create method")
-            if asyncio.iscoroutinefunction(provider.create):
-                response = await provider.create(prompt)
-            else:
-                response = provider.create(prompt)
-
-            if isinstance(response, ImageResponse):
-                return self._process_image_response(response)
-            elif isinstance(response, str):
-                image_response = ImageResponse([response], prompt)
-                return self._process_image_response(image_response)
+                raise ProviderNotFoundError(f'Provider not found: {provider}')
         else:
-            raise ValueError(f"Provider {provider} does not support image generation")
-
-        logging.error(f"Unexpected response type: {type(response)}")
-        raise NoImageResponseError(f"Unexpected response type: {type(response)}")
-
-    def _process_image_response(self, response: ImageResponse) -> ImagesResponse:
-        processed_images = []
-        for image_data in response.get_list():
-            if image_data.startswith('http://') or image_data.startswith('https://'):
-                processed_images.append(Image(url=image_data))
-            else:
-                image = to_image(image_data)
-                file_name = self._save_image(image)
-                processed_images.append(Image(url=file_name))
-        return ImagesResponse(processed_images)
-
-    def _save_image(self, image: 'PILImage') -> str:
-        os.makedirs('generated_images', exist_ok=True)
-        file_name = f"generated_images/image_{int(time.time())}.png"
-        image.save(file_name)
-        return file_name
-
-    async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs) -> ImagesResponse:
-        provider = self.models.get(model, self.provider)
-        if provider is None:
-            raise ValueError(f"Unknown model: {model}")
+            provider = self.models.get(model, self.provider)
+        return provider
 
-        if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
-            messages = [{"role": "user", "content": "create a variation of this image"}]
-            image_data = to_data_uri(image)
-            async for response in provider.create_async_generator(model, messages, image=image_data, **kwargs):
-                if isinstance(response, ImageResponse):
-                    return self._process_image_response(response)
-                elif isinstance(response, str):
-                    image_response = ImageResponse([response], "Image variation")
-                    return self._process_image_response(image_response)
-        elif hasattr(provider, 'create_variation'):
-            if asyncio.iscoroutinefunction(provider.create_variation):
-                response = await provider.create_variation(image, **kwargs)
-            else:
-                response = provider.create_variation(image, **kwargs)
-
-            if isinstance(response, ImageResponse):
-                return self._process_image_response(response)
-            elif isinstance(response, str):
-                image_response = ImageResponse([response], "Image variation")
-                return self._process_image_response(image_response)
+    async def generate(
+        self,
+        prompt,
+        model: str = "",
+        provider: ProviderType = None,
+        response_format: str = None,
+        connector: BaseConnector = None,
+        proxy: str = None,
+        **kwargs
+    ) -> ImagesResponse:
+        provider = self.get_provider(model, provider)
+        if hasattr(provider, "create_async_generator"):
+            response = create_image(
+                provider,
+                prompt,
+                **filter_none(
+                    response_format=response_format,
+                    connector=connector,
+                    proxy=self.client.get_proxy() if proxy is None else proxy,
+                ),
+                **kwargs
+            )
         else:
-            raise ValueError(f"Provider {provider} does not support image variation")
+            response = await provider.create_async(prompt)
+            return ImagesResponse([Image(image) for image in response.get_list()])
+        image = await iter_image_response(response, response_format, connector, proxy)
+        if image is None:
+            raise NoImageResponseError()
+        return image
 
-        raise NoImageResponseError("Failed to create image variation")
+    async def create_variation(
+        self,
+        image: ImageType,
+        model: str = None,
+        response_format: str = None,
+        connector: BaseConnector = None,
+        proxy: str = None,
+        **kwargs
+    ):
+        provider = self.get_provider(model, provider)
+        result = None
+        if hasattr(provider, "create_async_generator"):
+            response = provider.create_async_generator(
+                "",
+                [{"role": "user", "content": "create a image like this"}],
+                stream=True,
+                image=image,
+                **filter_none(
+                    response_format=response_format,
+                    connector=connector,
+                    proxy=self.client.get_proxy() if proxy is None else proxy,
+                ),
+                **kwargs
+            )
+            result = iter_image_response(response, response_format, connector, proxy)
+        if result is None:
+            raise NoImageResponseError()
+        return result