0.3.5.4

1 files changed, 180 insertions, 144 deletions

diff --git a/g4f/client/__init__.py b/g4f/client/__init__.py
index 549a244b..dea19a60 100644
--- a/g4f/client/__init__.py
+++ b/g4f/client/__init__.py
@@ -6,24 +6,27 @@ import random
 import string
 import asyncio
 import base64
-import aiohttp
 import logging
-from typing import Union, AsyncIterator, Iterator, Coroutine
+from typing import Union, AsyncIterator, Iterator, Coroutine, Optional
 
 from ..providers.base_provider import AsyncGeneratorProvider
-from ..image import ImageResponse, to_image, to_data_uri, is_accepted_format, EXTENSIONS_MAP
-from ..typing import Messages, Image
-from ..providers.types import ProviderType, FinishReason, BaseConversation
-from ..errors import NoImageResponseError
+from ..image import ImageResponse, copy_images, images_dir
+from ..typing import Messages, Image, ImageType
+from ..providers.types import ProviderType
+from ..providers.response import ResponseType, FinishReason, BaseConversation
+from ..errors import NoImageResponseError, ModelNotFoundError
 from ..providers.retry_provider import IterListProvider
+from ..providers.base_provider import get_running_loop
 from ..Provider.needs_auth.BingCreateImages import BingCreateImages
-from ..requests.aiohttp import get_connector
 from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse
 from .image_models import ImageModels
 from .types import IterResponse, ImageProvider, Client as BaseClient
 from .service import get_model_and_provider, get_last_provider, convert_to_provider
 from .helper import find_stop, filter_json, filter_none, safe_aclose, to_sync_iter, to_async_iterator
 
+ChatCompletionResponseType = Iterator[Union[ChatCompletion, ChatCompletionChunk, BaseConversation]]
+AsyncChatCompletionResponseType = AsyncIterator[Union[ChatCompletion, ChatCompletionChunk, BaseConversation]]
+
 try:
     anext # Python 3.8+
 except NameError:
@@ -35,12 +38,12 @@ except NameError:
 
 # Synchronous iter_response function
 def iter_response(
-    response: Union[Iterator[str], AsyncIterator[str]],
+    response: Union[Iterator[Union[str, ResponseType]]],
     stream: bool,
-    response_format: dict = None,
-    max_tokens: int = None,
-    stop: list = None
-) -> Iterator[Union[ChatCompletion, ChatCompletionChunk]]:
+    response_format: Optional[dict] = None,
+    max_tokens: Optional[int] = None,
+    stop: Optional[list[str]] = None
+) -> ChatCompletionResponseType:
     content = ""
     finish_reason = None
     completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
@@ -88,22 +91,23 @@ def iter_response(
         yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
 
 # Synchronous iter_append_model_and_provider function
-def iter_append_model_and_provider(response: Iterator[ChatCompletionChunk]) -> Iterator[ChatCompletionChunk]:
+def iter_append_model_and_provider(response: ChatCompletionResponseType) -> ChatCompletionResponseType:
     last_provider = None
 
     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")
-        yield chunk
+        if isinstance(chunk, (ChatCompletion, ChatCompletionChunk)):
+            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")
+            yield chunk
 
 async def async_iter_response(
-    response: AsyncIterator[str],
+    response: AsyncIterator[Union[str, ResponseType]],
     stream: bool,
-    response_format: dict = None,
-    max_tokens: int = None,
-    stop: list = None
-) -> AsyncIterator[Union[ChatCompletion, ChatCompletionChunk]]:
+    response_format: Optional[dict] = None,
+    max_tokens: Optional[int] = None,
+    stop: Optional[list[str]] = None
+) -> AsyncChatCompletionResponseType:
     content = ""
     finish_reason = None
     completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
@@ -149,13 +153,16 @@ async def async_iter_response(
         if hasattr(response, 'aclose'):
             await safe_aclose(response)
 
-async def async_iter_append_model_and_provider(response: AsyncIterator[ChatCompletionChunk]) -> AsyncIterator:
+async def async_iter_append_model_and_provider(
+        response: AsyncChatCompletionResponseType
+    ) -> AsyncChatCompletionResponseType:
     last_provider = None
     try:
         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")
+            if isinstance(chunk, (ChatCompletion, ChatCompletionChunk)):
+                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")
             yield chunk
     finally:
         if hasattr(response, 'aclose'):
@@ -164,8 +171,8 @@ async def async_iter_append_model_and_provider(response: AsyncIterator[ChatCompl
 class Client(BaseClient):
     def __init__(
         self,
-        provider: ProviderType = None,
-        image_provider: ImageProvider = None,
+        provider: Optional[ProviderType] = None,
+        image_provider: Optional[ImageProvider] = None,
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
@@ -173,7 +180,7 @@ class Client(BaseClient):
         self.images: Images = Images(self, image_provider)
 
 class Completions:
-    def __init__(self, client: Client, provider: ProviderType = None):
+    def __init__(self, client: Client, provider: Optional[ProviderType] = None):
         self.client: Client = client
         self.provider: ProviderType = provider
 
@@ -181,16 +188,16 @@ class Completions:
         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,
-        ignore_working: bool = False,
-        ignore_stream: bool = False,
+        provider: Optional[ProviderType] = None,
+        stream: Optional[bool] = False,
+        proxy: Optional[str] = None,
+        response_format: Optional[dict] = None,
+        max_tokens: Optional[int] = None,
+        stop: Optional[Union[list[str], str]] = None,
+        api_key: Optional[str] = None,
+        ignored: Optional[list[str]] = None,
+        ignore_working: Optional[bool] = False,
+        ignore_stream: Optional[bool] = False,
         **kwargs
     ) -> IterResponse:
         model, provider = get_model_and_provider(
@@ -234,22 +241,38 @@ class Completions:
 class Chat:
     completions: Completions
 
-    def __init__(self, client: Client, provider: ProviderType = None):
+    def __init__(self, client: Client, provider: Optional[ProviderType] = None):
         self.completions = Completions(client, provider)
 
 class Images:
-    def __init__(self, client: Client, provider: ProviderType = None):
+    def __init__(self, client: Client, provider: Optional[ProviderType] = None):
         self.client: Client = client
-        self.provider: ProviderType = provider
+        self.provider: Optional[ProviderType] = provider
         self.models: ImageModels = ImageModels(client)
 
-    def generate(self, prompt: str, model: str = None, provider: ProviderType = None, response_format: str = "url", proxy: str = None, **kwargs) -> ImagesResponse:
+    def generate(
+        self,
+        prompt: str,
+        model: str = None,
+        provider: Optional[ProviderType] = None,
+        response_format: str = "url",
+        proxy: Optional[str] = None,
+        **kwargs
+    ) -> ImagesResponse:
         """
         Synchronous generate method that runs the async_generate method in an event loop.
         """
-        return asyncio.run(self.async_generate(prompt, model, provider, response_format=response_format, proxy=proxy, **kwargs))
+        return asyncio.run(self.async_generate(prompt, model, provider, response_format, proxy, **kwargs))
 
-    async def async_generate(self, prompt: str, model: str = None, provider: ProviderType = None, response_format: str = "url", proxy: str = None, **kwargs) -> ImagesResponse:
+    async def async_generate(
+        self,
+        prompt: str,
+        model: Optional[str] = None,
+        provider: Optional[ProviderType] = None,
+        response_format: Optional[str] = "url",
+        proxy: Optional[str] = None,
+        **kwargs
+    ) -> ImagesResponse:
         if provider is None:
             provider_handler = self.models.get(model, provider or self.provider or BingCreateImages)
         elif isinstance(provider, str):
@@ -257,97 +280,73 @@ class Images:
         else:
             provider_handler = provider
         if provider_handler is None:
-            raise ValueError(f"Unknown model: {model}")
-        if proxy is None:
-            proxy = self.client.proxy
-
+            raise ModelNotFoundError(f"Unknown model: {model}")
         if isinstance(provider_handler, IterListProvider):
             if provider_handler.providers:
                 provider_handler = provider_handler.providers[0]
             else:
-                raise ValueError(f"IterListProvider for model {model} has no providers")
+                raise ModelNotFoundError(f"IterListProvider for model {model} has no providers")
+        if proxy is None:
+            proxy = self.client.proxy
 
         response = None
-        if hasattr(provider_handler, "create_async_generator"):
-            messages = [{"role": "user", "content": prompt}]
-            async for item in provider_handler.create_async_generator(model, messages, **kwargs):
+        if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
+            messages = [{"role": "user", "content": f"Generate a image: {prompt}"}]
+            async for item in provider_handler.create_async_generator(model, messages, prompt=prompt, **kwargs):
                 if isinstance(item, ImageResponse):
                     response = item
                     break
-        elif hasattr(provider, 'create'):
+        elif hasattr(provider_handler, 'create'):
             if asyncio.iscoroutinefunction(provider_handler.create):
                 response = await provider_handler.create(prompt)
             else:
                 response = provider_handler.create(prompt) 
             if isinstance(response, str):
                 response = ImageResponse([response], prompt)
+        elif hasattr(provider_handler, "create_completion"):
+            get_running_loop(check_nested=True)
+            messages = [{"role": "user", "content": f"Generate a image: {prompt}"}]
+            for item in provider_handler.create_completion(model, messages, prompt=prompt, **kwargs):
+                if isinstance(item, ImageResponse):
+                    response = item
+                    break
         else:
             raise ValueError(f"Provider {provider} does not support image generation")
         if isinstance(response, ImageResponse):
-            return await self._process_image_response(response, response_format, proxy, model=model, provider=provider)
-
+            return await self._process_image_response(
+                response,
+                response_format,
+                proxy,
+                model,
+                getattr(provider_handler, "__name__", None)
+            )
         raise NoImageResponseError(f"Unexpected response type: {type(response)}")
 
-    async def _process_image_response(self, response: ImageResponse, response_format: str, proxy: str = None, model: str = None, provider: str = None) -> ImagesResponse:
-            async def process_image_item(session: aiohttp.ClientSession, image_data: str):
-                image_data_bytes = None
-                if image_data.startswith("http://") or image_data.startswith("https://"):
-                    if response_format == "url":
-                        return Image(url=image_data, revised_prompt=response.alt)
-                    elif response_format == "b64_json":
-                        # Fetch the image data and convert it to base64
-                        image_data_bytes = await self._fetch_image(session, image_data)
-                        b64_json = base64.b64encode(image_data_bytes).decode("utf-8")
-                        return Image(b64_json=b64_json, url=image_data, revised_prompt=response.alt)
-                else:
-                    # Assume image_data is base64 data or binary
-                    if response_format == "url":
-                        if image_data.startswith("data:image"):
-                            # Remove the data URL scheme and get the base64 data
-                            base64_data = image_data.split(",", 1)[-1]
-                        else:
-                            base64_data = image_data
-                        # Decode the base64 data
-                        image_data_bytes = base64.b64decode(base64_data)
-                if image_data_bytes:
-                    file_name = self._save_image(image_data_bytes)
-                    return Image(url=file_name, revised_prompt=response.alt)
-                else:
-                    raise ValueError("Unable to process image data")
-
-            last_provider = get_last_provider(True)
-            async with aiohttp.ClientSession(cookies=response.get("cookies"), connector=get_connector(proxy=proxy)) as session:
-                return ImagesResponse(
-                    await asyncio.gather(*[process_image_item(session, image_data) for image_data in response.get_list()]),
-                    model=last_provider.get("model") if model is None else model,
-                    provider=last_provider.get("name") if provider is None else provider
-                )
-
-    async def _fetch_image(self, session: aiohttp.ClientSession, url: str) -> bytes:
-        # Asynchronously fetch image data from the URL
-        async with session.get(url) as resp:
-            if resp.status == 200:
-                return await resp.read()
-            else:
-                raise RuntimeError(f"Failed to fetch image from {url}, status code {resp.status}")
-
-    def _save_image(self, image_data_bytes: bytes) -> str:
-        os.makedirs('generated_images', exist_ok=True)
-        image = to_image(image_data_bytes)
-        file_name = f"generated_images/image_{int(time.time())}_{random.randint(0, 10000)}.{EXTENSIONS_MAP[is_accepted_format(image_data_bytes)]}"
-        image.save(file_name)
-        return file_name
-
-    def create_variation(self, image: Union[str, bytes], model: str = None, provider: ProviderType = None, response_format: str = "url", **kwargs) -> ImagesResponse:
+    def create_variation(
+        self,
+        image: Union[str, bytes],
+        model: str = None,
+        provider: Optional[ProviderType] = None,
+        response_format: str = "url",
+        **kwargs
+    ) -> ImagesResponse:
         return asyncio.run(self.async_create_variation(
            image, model, provider, response_format, **kwargs
         ))
 
-    async def async_create_variation(self, image: Union[str, bytes], model: str = None, provider: ProviderType = None, response_format: str = "url", proxy: str = None, **kwargs) -> ImagesResponse:
+    async def async_create_variation(
+        self,
+        image: ImageType,
+        model: Optional[str] = None,
+        provider: Optional[ProviderType] = None,
+        response_format: str = "url",
+        proxy: Optional[str] = None,
+        **kwargs
+    ) -> ImagesResponse:
         if provider is None:
             provider = self.models.get(model, provider or self.provider or BingCreateImages)
             if provider is None:
-                raise ValueError(f"Unknown model: {model}")
+                raise ModelNotFoundError(f"Unknown model: {model}")
         if isinstance(provider, str):
             provider = convert_to_provider(provider)
         if proxy is None:
@@ -355,38 +354,61 @@ class Images:
 
         if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
             messages = [{"role": "user", "content": "create a variation of this image"}]
-            image_data = to_data_uri(image)
             generator = None
             try:
-                generator = provider.create_async_generator(model, messages, image=image_data, response_format=response_format, proxy=proxy, **kwargs)
-                async for response in generator:
-                    if isinstance(response, ImageResponse):
-                        return self._process_image_response(response)
-            except RuntimeError as e:
-                if "async generator ignored GeneratorExit" in str(e):
-                    logging.warning("Generator ignored GeneratorExit in create_variation, handling gracefully")
-                else:
-                    raise
+                generator = provider.create_async_generator(model, messages, image=image, response_format=response_format, proxy=proxy, **kwargs)
+                async for chunk in generator:
+                    if isinstance(chunk, ImageResponse):
+                        response = chunk
+                        break
             finally:
                 if generator and hasattr(generator, 'aclose'):
                     await safe_aclose(generator)
-                logging.info("AsyncGeneratorProvider processing completed in create_variation")
         elif hasattr(provider, 'create_variation'):
             if asyncio.iscoroutinefunction(provider.create_variation):
                 response = await provider.create_variation(image, model=model, response_format=response_format, proxy=proxy, **kwargs)
             else:
                 response = provider.create_variation(image, model=model, response_format=response_format, proxy=proxy, **kwargs)
-            if isinstance(response, str):
-                response = ImageResponse([response])
-            return self._process_image_response(response)
         else:
-            raise ValueError(f"Provider {provider} does not support image variation")
+            raise NoImageResponseError(f"Provider {provider} does not support image variation")
+    
+        if isinstance(response, str):
+            response = ImageResponse([response])
+        if isinstance(response, ImageResponse):
+            return self._process_image_response(response, response_format, proxy, model, getattr(provider, "__name__", None))
+        raise NoImageResponseError(f"Unexpected response type: {type(response)}")
+
+    async def _process_image_response(
+        self,
+        response: ImageResponse,
+        response_format: str,
+        proxy: str = None,
+        model: Optional[str] = None,
+        provider: Optional[str] = None
+    ) -> list[Image]:
+        if response_format in ("url", "b64_json"):
+            images = await copy_images(response.get_list(), response.options.get("cookies"), proxy)
+            async def process_image_item(image_file: str) -> Image:
+                if response_format == "b64_json":
+                    with open(os.path.join(images_dir, os.path.basename(image_file)), "rb") as file:
+                        image_data = base64.b64encode(file.read()).decode()
+                        return Image(url=image_file, b64_json=image_data, revised_prompt=response.alt)
+                return Image(url=image_file, revised_prompt=response.alt)
+            images = await asyncio.gather(*[process_image_item(image) for image in images])
+        else:
+            images = [Image(url=image, revised_prompt=response.alt) for image in response.get_list()]
+        last_provider = get_last_provider(True)
+        return ImagesResponse(
+            images,
+            model=last_provider.get("model") if model is None else model,
+            provider=last_provider.get("name") if provider is None else provider
+        )
 
 class AsyncClient(BaseClient):
     def __init__(
         self,
-        provider: ProviderType = None,
-        image_provider: ImageProvider = None,
+        provider: Optional[ProviderType] = None,
+        image_provider: Optional[ImageProvider] = None,
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
@@ -396,11 +418,11 @@ class AsyncClient(BaseClient):
 class AsyncChat:
     completions: AsyncCompletions
 
-    def __init__(self, client: AsyncClient, provider: ProviderType = None):
+    def __init__(self, client: AsyncClient, provider: Optional[ProviderType] = None):
         self.completions = AsyncCompletions(client, provider)
 
 class AsyncCompletions:
-    def __init__(self, client: AsyncClient, provider: ProviderType = None):
+    def __init__(self, client: AsyncClient, provider: Optional[ProviderType] = None):
         self.client: AsyncClient = client
         self.provider: ProviderType = provider
 
@@ -408,18 +430,18 @@ class AsyncCompletions:
         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,
-        ignore_working: bool = False,
-        ignore_stream: bool = False,
+        provider: Optional[ProviderType] = None,
+        stream: Optional[bool] = False,
+        proxy: Optional[str] = None,
+        response_format: Optional[dict] = None,
+        max_tokens: Optional[int] = None,
+        stop: Optional[Union[list[str], str]] = None,
+        api_key: Optional[str] = None,
+        ignored: Optional[list[str]] = None,
+        ignore_working: Optional[bool] = False,
+        ignore_stream: Optional[bool] = False,
         **kwargs
-    ) -> Union[Coroutine[ChatCompletion], AsyncIterator[ChatCompletionChunk]]:
+    ) -> Union[Coroutine[ChatCompletion], AsyncIterator[ChatCompletionChunk, BaseConversation]]:
         model, provider = get_model_and_provider(
             model,
             self.provider if provider is None else provider,
@@ -450,15 +472,29 @@ class AsyncCompletions:
         return response if stream else anext(response)
 
 class AsyncImages(Images):
-    def __init__(self, client: AsyncClient, provider: ImageProvider = None):
+    def __init__(self, client: AsyncClient, provider: Optional[ProviderType] = None):
         self.client: AsyncClient = client
-        self.provider: ImageProvider = provider
+        self.provider: Optional[ProviderType] = provider
         self.models: ImageModels = ImageModels(client)
 
-    async def generate(self, prompt: str, model: str = None, provider: ProviderType = None, response_format: str = "url", **kwargs) -> ImagesResponse:
+    async def generate(
+        self,
+        prompt: str,
+        model: Optional[str] = None,
+        provider: Optional[ProviderType] = None,
+        response_format: str = "url",
+        **kwargs
+    ) -> ImagesResponse:
         return await self.async_generate(prompt, model, provider, response_format, **kwargs)
 
-    async def create_variation(self, image: Union[str, bytes], model: str = None, provider: ProviderType = None, response_format: str = "url", **kwargs) -> ImagesResponse:
+    async def create_variation(
+        self,
+        image: ImageType,
+        model: str = None,
+        provider: ProviderType = None,
+        response_format: str = "url",
+        **kwargs
+    ) -> ImagesResponse:
         return await self.async_create_variation(
            image, model, provider, response_format, **kwargs
-        )
+        )
\ No newline at end of file