summaryrefslogtreecommitdiffstats
path: root/g4f/client/async_client.py
diff options
context:
space:
mode:
Diffstat (limited to 'g4f/client/async_client.py')
-rw-r--r--g4f/client/async_client.py322
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