1 files changed, 20 insertions, 94 deletions

diff --git a/g4f/client.py b/g4f/client/client.py
index 2c4fe788..8ce5d932 100644
--- a/g4f/client.py
+++ b/g4f/client/client.py
@@ -1,40 +1,19 @@
 from __future__ import annotations
 
-import re
-import os
 import time
 import random
 import string
 
+from ..typing import Union, Iterator, Messages, ImageType
+from ..providers.types import BaseProvider, ProviderType, FinishReason
+from ..image import ImageResponse as ImageProviderResponse
+from ..errors import NoImageResponseError
 from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse
-from .typing import Union, Iterator, Messages, ImageType
-from .providers.types import BaseProvider, ProviderType, FinishReason
-from .image import ImageResponse as ImageProviderResponse
-from .errors import NoImageResponseError, RateLimitError, MissingAuthError
-from . import get_model_and_provider, get_last_provider
-
-from .Provider.BingCreateImages import BingCreateImages
-from .Provider.needs_auth import Gemini, OpenaiChat
-from .Provider.You import You
-
-ImageProvider = Union[BaseProvider, object]
-Proxies = Union[dict, str]
-IterResponse = Iterator[Union[ChatCompletion, ChatCompletionChunk]]
-
-def read_json(text: str) -> dict:
-    """
-    Parses JSON code block from a string.
-
-    Args:
-        text (str): A string containing a JSON code block.
-
-    Returns:
-        dict: A dictionary parsed from the JSON code block.
-    """
-    match = re.search(r"```(json|)\n(?P<code>[\S\s]+?)\n```", text)
-    if match:
-        return match.group("code")
-    return text
+from .image_models import ImageModels
+from .types import IterResponse, ImageProvider
+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
 
 def iter_response(
     response: iter[str],
@@ -53,20 +32,7 @@ def iter_response(
         content += str(chunk)
         if max_tokens is not None and idx + 1 >= max_tokens:
             finish_reason = "length"
-        first = -1
-        word = None
-        if stop is not None:
-            for word in list(stop):
-                first = content.find(word)
-                if first != -1:
-                    content = content[:first]
-                    break
-            if stream and first != -1:
-                first = chunk.find(word)
-                if first != -1:
-                    chunk = chunk[:first]
-                else:
-                    first = 0
+        first, content, chunk = find_stop(stop, content, chunk if stream else None)
         if first != -1:
             finish_reason = "stop"
         if stream:
@@ -79,7 +45,7 @@ def iter_response(
     else:
         if response_format is not None and "type" in response_format:
             if response_format["type"] == "json_object":
-                content = read_json(content)
+                content = filter_json(content)
         yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
 
 def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
@@ -90,37 +56,17 @@ def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
         chunk.provider =  last_provider.get("name")
         yield chunk
 
-class Client():
-
+class Client(BaseClient):
     def __init__(
         self,
-        api_key: str = None,
-        proxies: Proxies = None,
         provider: ProviderType = None,
         image_provider: ImageProvider = None,
         **kwargs
     ) -> None:
-        self.api_key: str = api_key
-        self.proxies: Proxies = proxies
+        super().__init__(**kwargs)
         self.chat: Chat = Chat(self, provider)
         self.images: Images = Images(self, image_provider)
 
-    def get_proxy(self) -> Union[str, None]:
-        if isinstance(self.proxies, str):
-            return self.proxies
-        elif self.proxies is None:
-            return os.environ.get("G4F_PROXY")
-        elif "all" in self.proxies:
-            return self.proxies["all"]
-        elif "https" in self.proxies:
-            return self.proxies["https"]
-
-def filter_none(**kwargs):
-    for key in list(kwargs.keys()):
-        if kwargs[key] is None:
-            del kwargs[key]
-    return kwargs
-
 class Completions():
     def __init__(self, client: Client, provider: ProviderType = None):
         self.client: Client = client
@@ -132,6 +78,7 @@ class Completions():
         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,
@@ -148,13 +95,12 @@ class Completions():
             ignored,
             ignore_working,
             ignore_stream,
-            **kwargs
         )
         stop = [stop] if isinstance(stop, str) else stop
         response = provider.create_completion(
             model, messages, stream,            
             **filter_none(
-                proxy=self.client.get_proxy(),
+                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
@@ -171,18 +117,6 @@ class Chat():
     def __init__(self, client: Client, provider: ProviderType = None):
         self.completions = Completions(client, provider)
 
-class ImageModels():
-    gemini = Gemini
-    openai = OpenaiChat
-    you = You
-
-    def __init__(self, client: Client) -> None:
-        self.client = client
-        self.default = BingCreateImages(proxy=self.client.get_proxy())
-
-    def get(self, name: str, default: ImageProvider = None) -> ImageProvider:
-        return getattr(self, name) if hasattr(self, name) else default or self.default
-
 def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
     for chunk in list(response):
         if isinstance(chunk, ImageProviderResponse):
@@ -190,10 +124,12 @@ def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
 
 def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator:
     prompt = f"create a image with: {prompt}"
+    if provider.__name__ == "You":
+        kwargs["chat_mode"] = "create"
     return provider.create_completion(
         model,
         [{"role": "user", "content": prompt}],
-        True,
+        stream=True,
         proxy=client.get_proxy(),
         **kwargs
     )
@@ -209,14 +145,7 @@ class Images():
         if isinstance(provider, type) and issubclass(provider, BaseProvider):
             response = create_image(self.client, provider, prompt, **kwargs)
         else:
-            try:
-                response = list(provider.create(prompt))
-            except (RateLimitError, MissingAuthError) as e:
-                # Fallback for default provider
-                if self.provider is None:
-                    response = create_image(self.client, self.models.you, prompt, model or "dall-e", **kwargs)
-                else:
-                    raise e
+            response = list(provider.create(prompt))
         image = iter_image_response(response)
         if image is None:
             raise NoImageResponseError()
@@ -234,10 +163,7 @@ class Images():
                 proxy=self.client.get_proxy(),
                 **kwargs
             )
-            for chunk in response:
-                if isinstance(chunk, ImageProviderResponse):
-                    result = ([chunk.images] if isinstance(chunk.images, str) else chunk.images)
-                    result = ImagesResponse([Image(image)for image in result])
+            result = iter_image_response(response)
         if result is None:
             raise NoImageResponseError()
         return result
\ No newline at end of file