diff options
Diffstat (limited to 'g4f/client/client.py')
| -rw-r--r-- | g4f/client/client.py | 211 |
1 files changed, 170 insertions, 41 deletions
diff --git a/g4f/client/client.py b/g4f/client/client.py index 63bae4fe..56644913 100644 --- a/g4f/client/client.py +++ b/g4f/client/client.py @@ -1,9 +1,14 @@ from __future__ import annotations +import os import time import random import string - +import logging +import asyncio +from typing import Union +from ..providers.base_provider import AsyncGeneratorProvider +from ..image import ImageResponse, to_image, to_data_uri from ..typing import Union, Iterator, Messages, ImageType from ..providers.types import BaseProvider, ProviderType, FinishReason from ..providers.conversation import BaseConversation @@ -15,9 +20,12 @@ 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 +from ..models import ModelUtils +from ..Provider import IterListProvider + def iter_response( - response: iter[str], + response: Iterator[str], stream: bool, response_format: dict = None, max_tokens: int = None, @@ -26,6 +34,7 @@ def iter_response( content = "" finish_reason = None completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28)) + for idx, chunk in enumerate(response): if isinstance(chunk, FinishReason): finish_reason = chunk.reason @@ -33,17 +42,25 @@ def iter_response( elif isinstance(chunk, BaseConversation): yield chunk continue + content += str(chunk) + if max_tokens is not None and idx + 1 >= max_tokens: finish_reason = "length" + first, content, chunk = find_stop(stop, content, chunk if stream else None) + 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: @@ -52,14 +69,16 @@ def iter_response( content = filter_json(content) yield ChatCompletion(content, finish_reason, completion_id, int(time.time())) + def iter_append_model_and_provider(response: IterResponse) -> IterResponse: 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") + chunk.provider = last_provider.get("name") yield chunk + class Client(BaseClient): def __init__( self, @@ -69,9 +88,17 @@ class Client(BaseClient): ) -> None: super().__init__(**kwargs) self.chat: Chat = Chat(self, provider) - self.images: Images = Images(self, image_provider) + self._images: Images = Images(self, image_provider) -class Completions(): + @property + def images(self) -> Images: + return self._images + + async def async_images(self) -> Images: + return self._images + + +class Completions: def __init__(self, client: Client, provider: ProviderType = None): self.client: Client = client self.provider: ProviderType = provider @@ -87,7 +114,7 @@ class Completions(): max_tokens: int = None, stop: Union[list[str], str] = None, api_key: str = None, - ignored : list[str] = None, + ignored: list[str] = None, ignore_working: bool = False, ignore_stream: bool = False, **kwargs @@ -100,11 +127,13 @@ class Completions(): ignore_working, ignore_stream, ) - + stop = [stop] if isinstance(stop, str) else stop + response = provider.create_completion( - model, messages, - stream=stream, + model, + messages, + stream=stream, **filter_none( proxy=self.client.get_proxy() if proxy is None else proxy, max_tokens=max_tokens, @@ -113,66 +142,166 @@ class Completions(): ), **kwargs ) + response = iter_response(response, stream, response_format, max_tokens, stop) response = iter_append_model_and_provider(response) + return response if stream else next(response) -class Chat(): + +class Chat: completions: Completions def __init__(self, client: Client, provider: ProviderType = None): self.completions = Completions(client, provider) + def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]: - for chunk in list(response): + logging.info("Starting iter_image_response") + response_list = list(response) + logging.info(f"Response list: {response_list}") + + for chunk in response_list: + logging.info(f"Processing chunk: {chunk}") if isinstance(chunk, ImageProviderResponse): + logging.info("Found ImageProviderResponse") return ImagesResponse([Image(image) for image in chunk.get_list()]) - -def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator: + + logging.warning("No ImageProviderResponse found in the response") + return None +def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator: + logging.info(f"Creating image with provider: {provider}, model: {model}, prompt: {prompt}") + if isinstance(provider, type) and provider.__name__ == "You": kwargs["chat_mode"] = "create" else: - prompt = f"create a image with: {prompt}" - return provider.create_completion( + prompt = f"create an image with: {prompt}" + + response = provider.create_completion( model, [{"role": "user", "content": prompt}], stream=True, proxy=client.get_proxy(), **kwargs ) + + logging.info(f"Response from create_completion: {response}") + return response -class Images(): - def __init__(self, client: Client, provider: ImageProvider = None): - self.client: Client = client + +class Images: + def __init__(self, client: 'Client', provider: ImageProvider = None): + self.client: 'Client' = client self.provider: ImageProvider = provider self.models: ImageModels = ImageModels(client) - def generate(self, prompt, model: str = None, **kwargs) -> ImagesResponse: + def generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse: + logging.info(f"Starting synchronous image generation for model: {model}, prompt: {prompt}") + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(self.async_generate(prompt, model, **kwargs)) + logging.info(f"Synchronous image generation completed. Result: {result}") + return result + except Exception as e: + logging.error(f"Error in synchronous image generation: {str(e)}") + raise + finally: + if loop.is_running(): + loop.close() + + async def async_generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse: + logging.info(f"Generating image for model: {model}, prompt: {prompt}") provider = self.models.get(model, self.provider) - if isinstance(provider, type) and issubclass(provider, BaseProvider): - response = create_image(self.client, provider, prompt, **kwargs) + 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}") + 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) else: - response = list(provider.create(prompt)) - image = iter_image_response(response) - if image is None: - raise NoImageResponseError() - return image + 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 - def create_variation(self, image: ImageType, model: str = None, **kwargs): + async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs): provider = self.models.get(model, self.provider) - result = None - if isinstance(provider, type) and issubclass(provider, BaseProvider): - response = provider.create_completion( - "", - [{"role": "user", "content": "create a image like this"}], - True, - image=image, - proxy=self.client.get_proxy(), - **kwargs - ) - result = iter_image_response(response) - if result is None: - raise NoImageResponseError() - return result
\ No newline at end of file + if provider is None: + raise ValueError(f"Unknown model: {model}") + + 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) + else: + raise ValueError(f"Provider {provider} does not support image variation") + + raise NoImageResponseError("Failed to create image variation") + |