diff options
Diffstat (limited to '')
| -rw-r--r-- | g4f/client/__init__.py | 324 |
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 |