summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorkqlio67 <kqlio67@users.noreply.github.com>2024-10-15 11:22:52 +0200
committerkqlio67 <kqlio67@users.noreply.github.com>2024-10-15 11:22:52 +0200
commit0d868f64aa5fa4bc3891737bc76e82843e4bb5b0 (patch)
tree35dd18186194de23f999738e56dfdb9c5adae764
parentrefactor(etc/unittest): update async client test structure (diff)
downloadgpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar.gz
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar.bz2
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar.lz
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar.xz
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.tar.zst
gpt4free-0d868f64aa5fa4bc3891737bc76e82843e4bb5b0.zip
-rw-r--r--g4f/client/__init__.py1
-rw-r--r--g4f/client/async_client.py339
-rw-r--r--g4f/client/client.py467
3 files changed, 335 insertions, 472 deletions
diff --git a/g4f/client/__init__.py b/g4f/client/__init__.py
index 5bb4ba35..9fb3551e 100644
--- a/g4f/client/__init__.py
+++ b/g4f/client/__init__.py
@@ -1,3 +1,2 @@
from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse
from .client import Client
-from .async_client import AsyncClient \ No newline at end of file
diff --git a/g4f/client/async_client.py b/g4f/client/async_client.py
deleted file mode 100644
index b4d52a60..00000000
--- a/g4f/client/async_client.py
+++ /dev/null
@@ -1,339 +0,0 @@
-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
-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 .helper import cast_iter_async
-
-try:
- anext # Python 3.8+
-except NameError:
- async def anext(aiter):
- try:
- return await aiter.__anext__()
- except StopAsyncIteration:
- raise StopIteration
-
-async def safe_aclose(generator):
- try:
- await generator.aclose()
- except Exception as e:
- logging.warning(f"Error while closing generator: {e}")
-
-async def iter_response(
- response: AsyncIterator[str],
- stream: bool,
- response_format: dict = None,
- max_tokens: int = None,
- stop: list = None
-) -> AsyncIterator[Union[ChatCompletion, ChatCompletionChunk]]:
- content = ""
- finish_reason = None
- completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
- idx = 0
-
- try:
- async for chunk in response:
- if isinstance(chunk, FinishReason):
- finish_reason = chunk.reason
- break
- elif isinstance(chunk, BaseConversation):
- yield chunk
- continue
-
- content += str(chunk)
- idx += 1
-
- if max_tokens is not None and idx >= 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:
- if response_format is not None and "type" in response_format:
- if response_format["type"] == "json_object":
- content = filter_json(content)
- yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
- finally:
- if hasattr(response, 'aclose'):
- await safe_aclose(response)
-
-async def iter_append_model_and_provider(response: AsyncIterator) -> AsyncIterator:
- 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")
- yield chunk
- finally:
- if hasattr(response, 'aclose'):
- await safe_aclose(response)
-
-class AsyncClient(BaseClient):
- def __init__(
- self,
- 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)
-
- @property
- def images(self) -> Images:
- return self._images
-
-class Completions:
- def __init__(self, client: 'AsyncClient', provider: ProviderType = None):
- self.client: 'AsyncClient' = client
- self.provider: ProviderType = provider
-
- async 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,
- ignore_working: bool = False,
- ignore_stream: bool = False,
- **kwargs
- ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
- model, provider = get_model_and_provider(
- model,
- self.provider if provider is None else provider,
- stream,
- ignored,
- ignore_working,
- ignore_stream,
- )
-
- stop = [stop] if isinstance(stop, str) else stop
-
- response = 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
- ),
- **kwargs
- )
-
- if isinstance(response, AsyncIterator):
- 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)
- else:
- response = cast_iter_async(response)
- 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)
-
-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")
- try:
- 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
- finally:
- if hasattr(response, 'aclose'):
- await safe_aclose(response)
-
-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 isinstance(provider, type) and provider.__name__ == "You":
- kwargs["chat_mode"] = "create"
- else:
- prompt = f"create an image with: {prompt}"
-
- response = await 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: '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}")
- 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}]
- generator = None
- try:
- generator = provider.create_async_generator(model, messages, **kwargs)
- async for response in generator:
- logging.debug(f"Received response: {type(response)}")
- 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)
- except RuntimeError as e:
- if "async generator ignored GeneratorExit" in str(e):
- logging.warning("Generator ignored GeneratorExit, handling gracefully")
- else:
- raise
- finally:
- if generator and hasattr(generator, 'aclose'):
- await safe_aclose(generator)
- logging.info("AsyncGeneratorProvider processing completed")
- elif hasattr(provider, 'create'):
- logging.info("Using provider's create method")
- async_create = asyncio.iscoroutinefunction(provider.create)
- if async_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)
- elif hasattr(provider, 'create_completion'):
- logging.info("Using provider's create_completion method")
- response = await create_image(self.client, provider, prompt, model, **kwargs)
- async for chunk in response:
- if isinstance(chunk, ImageProviderResponse):
- logging.info("Found ImageProviderResponse")
- return ImagesResponse([Image(image) for image in chunk.get_list()])
- 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}")
-
- 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, **kwargs)
- async for response in generator:
- 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)
- except RuntimeError as e:
- if "async generator ignored GeneratorExit" in str(e):
- logging.warning("Generator ignored GeneratorExit in create_variation, handling gracefully")
- else:
- raise
- 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, **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")
diff --git a/g4f/client/client.py b/g4f/client/client.py
index 56644913..41238df5 100644
--- a/g4f/client/client.py
+++ b/g4f/client/client.py
@@ -4,12 +4,16 @@ import os
import time
import random
import string
-import logging
+import threading
import asyncio
-from typing import Union
+import base64
+import aiohttp
+import queue
+from typing import Union, AsyncIterator, Iterator
+
from ..providers.base_provider import AsyncGeneratorProvider
from ..image import ImageResponse, to_image, to_data_uri
-from ..typing import Union, Iterator, Messages, ImageType
+from ..typing import Messages, ImageType
from ..providers.types import BaseProvider, ProviderType, FinishReason
from ..providers.conversation import BaseConversation
from ..image import ImageResponse as ImageProviderResponse
@@ -23,44 +27,83 @@ from .helper import find_stop, filter_json, filter_none
from ..models import ModelUtils
from ..Provider import IterListProvider
+# Helper function to convert an async generator to a synchronous iterator
+def to_sync_iter(async_gen: AsyncIterator) -> Iterator:
+ q = queue.Queue()
+ loop = asyncio.new_event_loop()
+ done = object()
+
+ def _run():
+ asyncio.set_event_loop(loop)
+
+ async def iterate():
+ try:
+ async for item in async_gen:
+ q.put(item)
+ finally:
+ q.put(done)
+
+ loop.run_until_complete(iterate())
+ loop.close()
+
+ threading.Thread(target=_run).start()
+ while True:
+ item = q.get()
+ if item is done:
+ break
+ yield item
+
+# Helper function to convert a synchronous iterator to an async iterator
+async def to_async_iterator(iterator):
+ for item in iterator:
+ yield item
+
+# Synchronous iter_response function
def iter_response(
- response: Iterator[str],
+ response: Union[Iterator[str], AsyncIterator[str]],
stream: bool,
response_format: dict = None,
max_tokens: int = None,
stop: list = None
-) -> IterResponse:
+) -> Iterator[Union[ChatCompletion, ChatCompletionChunk]]:
content = ""
finish_reason = None
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
-
- for idx, chunk in enumerate(response):
+ idx = 0
+
+ if hasattr(response, '__aiter__'):
+ # It's an async iterator, wrap it into a sync iterator
+ response = to_sync_iter(response)
+
+ 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:
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
-
+
+ idx += 1
+
finish_reason = "stop" if finish_reason is None else finish_reason
-
+
if stream:
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
else:
@@ -69,16 +112,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:
+# Synchronous iter_append_model_and_provider function
+def iter_append_model_and_provider(response: Iterator) -> Iterator:
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
-
class Client(BaseClient):
def __init__(
self,
@@ -97,7 +140,6 @@ class Client(BaseClient):
async def async_images(self) -> Images:
return self._images
-
class Completions:
def __init__(self, client: Client, provider: ProviderType = None):
self.client: Client = client
@@ -129,25 +171,115 @@ class Completions:
)
stop = [stop] if isinstance(stop, str) else stop
-
- response = provider.create_completion(
+
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ # Run the asynchronous function in an event loop
+ response = asyncio.run(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
+ ),
+ **kwargs
+ ))
+ else:
+ response = 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
+ ),
+ **kwargs
+ )
+
+ if stream:
+ if hasattr(response, '__aiter__'):
+ # It's an async generator, wrap it into a sync iterator
+ response = to_sync_iter(response)
+
+ # Now 'response' is an iterator
+ response = iter_response(response, stream, response_format, max_tokens, stop)
+ response = iter_append_model_and_provider(response)
+ return response
+ else:
+ if hasattr(response, '__aiter__'):
+ # If response is an async generator, collect it into a list
+ response = list(to_sync_iter(response))
+ response = iter_response(response, stream, response_format, max_tokens, stop)
+ response = iter_append_model_and_provider(response)
+ return next(response)
+
+ async def async_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,
+ ignore_working: bool = False,
+ ignore_stream: bool = False,
+ **kwargs
+ ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
+ model, provider = get_model_and_provider(
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
- ),
- **kwargs
+ self.provider if provider is None else provider,
+ stream,
+ ignored,
+ ignore_working,
+ ignore_stream,
)
-
- response = iter_response(response, stream, response_format, max_tokens, stop)
- response = iter_append_model_and_provider(response)
-
- return response if stream else next(response)
+ stop = [stop] if isinstance(stop, str) else stop
+
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ 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
+ ),
+ **kwargs
+ )
+ else:
+ response = 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
+ ),
+ **kwargs
+ )
+
+ # Removed 'await' here since 'async_iter_response' returns an async generator
+ response = async_iter_response(response, stream, response_format, max_tokens, stop)
+ response = async_iter_append_model_and_provider(response)
+
+ if stream:
+ return response
+ else:
+ async for result in response:
+ return result
class Chat:
completions: Completions
@@ -155,153 +287,224 @@ class Chat:
def __init__(self, client: Client, provider: ProviderType = None):
self.completions = Completions(client, provider)
+# Asynchronous versions of the helper functions
+async def async_iter_response(
+ response: Union[AsyncIterator[str], Iterator[str]],
+ stream: bool,
+ response_format: dict = None,
+ max_tokens: int = None,
+ stop: list = None
+) -> AsyncIterator[Union[ChatCompletion, ChatCompletionChunk]]:
+ content = ""
+ finish_reason = None
+ completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
+ idx = 0
+
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
+ 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:
+ finish_reason = "length"
+
+ first, content, chunk = find_stop(stop, content, chunk if stream else None)
-def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
- 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 first != -1:
+ finish_reason = "stop"
+
+ if stream:
+ yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
+
+ if finish_reason is not None:
+ break
+
+ idx += 1
+
+ finish_reason = "stop" if finish_reason is None else finish_reason
+
+ if stream:
+ yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
+ else:
+ if response_format is not None and "type" in response_format:
+ if response_format["type"] == "json_object":
+ content = filter_json(content)
+ yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
+
+async def async_iter_append_model_and_provider(response: AsyncIterator) -> AsyncIterator:
+ last_provider = None
+
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
+ 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")
+ yield chunk
+
+async def iter_image_response(response: AsyncIterator) -> Union[ImagesResponse, None]:
+ response_list = []
+ async for chunk in response:
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
+ response_list.extend(chunk.get_list())
+ elif isinstance(chunk, str):
+ response_list.append(chunk)
+ if response_list:
+ return ImagesResponse([Image(image) for image in response_list])
-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}")
-
+ return None
+
+async def create_image(client: Client, 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 = provider.create_completion(
- model,
- [{"role": "user", "content": prompt}],
- stream=True,
- proxy=client.get_proxy(),
- **kwargs
- )
-
- logging.info(f"Response from create_completion: {response}")
+
+ if asyncio.iscoroutinefunction(provider.create_completion):
+ response = await provider.create_completion(
+ model,
+ [{"role": "user", "content": prompt}],
+ stream=True,
+ proxy=client.get_proxy(),
+ **kwargs
+ )
+ else:
+ response = provider.create_completion(
+ model,
+ [{"role": "user", "content": prompt}],
+ stream=True,
+ proxy=client.get_proxy(),
+ **kwargs
+ )
+
+ # Wrap synchronous iterator into async iterator if necessary
+ if not hasattr(response, '__aiter__'):
+ response = to_async_iterator(response)
+
return response
+class Image:
+ def __init__(self, url: str = None, b64_json: str = None):
+ self.url = url
+ self.b64_json = b64_json
+
+ def __repr__(self):
+ return f"Image(url={self.url}, b64_json={'<base64 data>' if self.b64_json else None})"
+
+class ImagesResponse:
+ def __init__(self, data: list[Image]):
+ self.data = data
+
+ def __repr__(self):
+ return f"ImagesResponse(data={self.data})"
class Images:
- def __init__(self, client: 'Client', provider: ImageProvider = None):
+ def __init__(self, client: 'Client', provider: 'ImageProvider' = None):
self.client: 'Client' = client
- self.provider: ImageProvider = provider
+ self.provider: 'ImageProvider' = provider
self.models: ImageModels = ImageModels(client)
- 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}")
+ def generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
+ """
+ Synchronous generate method that runs the async_generate method in an event loop.
+ """
+ return asyncio.run(self.async_generate(prompt, model, response_format=response_format, **kwargs))
+
+ async def async_generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
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}")
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)
+ return await self._process_image_response(response, response_format)
elif isinstance(response, str):
image_response = ImageResponse([response], prompt)
- return self._process_image_response(image_response)
+ return await self._process_image_response(image_response, response_format)
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)
+ return await self._process_image_response(response, response_format)
elif isinstance(response, str):
image_response = ImageResponse([response], prompt)
- return self._process_image_response(image_response)
+ return await self._process_image_response(image_response, response_format)
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:
+ async def _process_image_response(self, response: ImageResponse, response_format: str) -> 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))
+ if response_format == "url":
+ processed_images.append(Image(url=image_data))
+ elif response_format == "b64_json":
+ # Fetch the image data and convert it to base64
+ image_content = await self._fetch_image(image_data)
+ b64_json = base64.b64encode(image_content).decode('utf-8')
+ processed_images.append(Image(b64_json=b64_json))
else:
- image = to_image(image_data)
- file_name = self._save_image(image)
- processed_images.append(Image(url=file_name))
+ # 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
+ header, base64_data = image_data.split(',', 1)
+ else:
+ base64_data = image_data
+ # Decode the base64 data
+ image_data_bytes = base64.b64decode(base64_data)
+ # Convert bytes to an image
+ image = to_image(image_data_bytes)
+ file_name = self._save_image(image)
+ processed_images.append(Image(url=file_name))
+ elif response_format == "b64_json":
+ if isinstance(image_data, bytes):
+ b64_json = base64.b64encode(image_data).decode('utf-8')
+ else:
+ b64_json = image_data # If already base64-encoded string
+ processed_images.append(Image(b64_json=b64_json))
+
return ImagesResponse(processed_images)
+ async def _fetch_image(self, url: str) -> bytes:
+ # Asynchronously fetch image data from the URL
+ async with aiohttp.ClientSession() as session:
+ async with session.get(url) as resp:
+ if resp.status == 200:
+ return await resp.read()
+ else:
+ raise Exception(f"Failed to fetch image from {url}, status code {resp.status}")
+
def _save_image(self, image: 'PILImage') -> str:
os.makedirs('generated_images', exist_ok=True)
- file_name = f"generated_images/image_{int(time.time())}.png"
+ file_name = f"generated_images/image_{int(time.time())}_{random.randint(0, 10000)}.png"
image.save(file_name)
return file_name
- async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs):
- provider = self.models.get(model, self.provider)
- 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")
-
+ async def create_variation(self, image: Union[str, bytes], model: str = None, response_format: str = "url", **kwargs):
+ # Existing implementation, adjust if you want to support b64_json here as well
+ pass