mirror of
https://github.com/xtekky/gpt4free.git
synced 2025-12-06 02:30:41 -08:00
feat(g4f/client/async_client.py, g4f/client/async_client.py): enhance async and sync handling in client
This commit is contained in:
parent
e1e37b34a5
commit
0d868f64aa
3 changed files with 336 additions and 473 deletions
|
|
@ -1,3 +1,2 @@
|
||||||
from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse
|
from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse
|
||||||
from .client import Client
|
from .client import Client
|
||||||
from .async_client import AsyncClient
|
|
||||||
|
|
@ -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")
|
|
||||||
|
|
@ -4,12 +4,16 @@ import os
|
||||||
import time
|
import time
|
||||||
import random
|
import random
|
||||||
import string
|
import string
|
||||||
import logging
|
import threading
|
||||||
import asyncio
|
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 ..providers.base_provider import AsyncGeneratorProvider
|
||||||
from ..image import ImageResponse, to_image, to_data_uri
|
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.types import BaseProvider, ProviderType, FinishReason
|
||||||
from ..providers.conversation import BaseConversation
|
from ..providers.conversation import BaseConversation
|
||||||
from ..image import ImageResponse as ImageProviderResponse
|
from ..image import ImageResponse as ImageProviderResponse
|
||||||
|
|
@ -23,44 +27,83 @@ from .helper import find_stop, filter_json, filter_none
|
||||||
from ..models import ModelUtils
|
from ..models import ModelUtils
|
||||||
from ..Provider import IterListProvider
|
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(
|
def iter_response(
|
||||||
response: Iterator[str],
|
response: Union[Iterator[str], AsyncIterator[str]],
|
||||||
stream: bool,
|
stream: bool,
|
||||||
response_format: dict = None,
|
response_format: dict = None,
|
||||||
max_tokens: int = None,
|
max_tokens: int = None,
|
||||||
stop: list = None
|
stop: list = None
|
||||||
) -> IterResponse:
|
) -> Iterator[Union[ChatCompletion, ChatCompletionChunk]]:
|
||||||
content = ""
|
content = ""
|
||||||
finish_reason = None
|
finish_reason = None
|
||||||
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
|
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
|
||||||
|
idx = 0
|
||||||
for idx, chunk in enumerate(response):
|
|
||||||
|
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):
|
if isinstance(chunk, FinishReason):
|
||||||
finish_reason = chunk.reason
|
finish_reason = chunk.reason
|
||||||
break
|
break
|
||||||
elif isinstance(chunk, BaseConversation):
|
elif isinstance(chunk, BaseConversation):
|
||||||
yield chunk
|
yield chunk
|
||||||
continue
|
continue
|
||||||
|
|
||||||
content += str(chunk)
|
content += str(chunk)
|
||||||
|
|
||||||
if max_tokens is not None and idx + 1 >= max_tokens:
|
if max_tokens is not None and idx + 1 >= max_tokens:
|
||||||
finish_reason = "length"
|
finish_reason = "length"
|
||||||
|
|
||||||
first, content, chunk = find_stop(stop, content, chunk if stream else None)
|
first, content, chunk = find_stop(stop, content, chunk if stream else None)
|
||||||
|
|
||||||
if first != -1:
|
if first != -1:
|
||||||
finish_reason = "stop"
|
finish_reason = "stop"
|
||||||
|
|
||||||
if stream:
|
if stream:
|
||||||
yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
|
yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
|
||||||
|
|
||||||
if finish_reason is not None:
|
if finish_reason is not None:
|
||||||
break
|
break
|
||||||
|
|
||||||
|
idx += 1
|
||||||
|
|
||||||
finish_reason = "stop" if finish_reason is None else finish_reason
|
finish_reason = "stop" if finish_reason is None else finish_reason
|
||||||
|
|
||||||
if stream:
|
if stream:
|
||||||
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
|
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
|
||||||
else:
|
else:
|
||||||
|
|
@ -69,16 +112,16 @@ def iter_response(
|
||||||
content = filter_json(content)
|
content = filter_json(content)
|
||||||
yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
|
yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
|
||||||
|
|
||||||
|
# Synchronous iter_append_model_and_provider function
|
||||||
def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
|
def iter_append_model_and_provider(response: Iterator) -> Iterator:
|
||||||
last_provider = None
|
last_provider = None
|
||||||
|
|
||||||
for chunk in response:
|
for chunk in response:
|
||||||
last_provider = get_last_provider(True) if last_provider is None else last_provider
|
last_provider = get_last_provider(True) if last_provider is None else last_provider
|
||||||
chunk.model = last_provider.get("model")
|
chunk.model = last_provider.get("model")
|
||||||
chunk.provider = last_provider.get("name")
|
chunk.provider = last_provider.get("name")
|
||||||
yield chunk
|
yield chunk
|
||||||
|
|
||||||
|
|
||||||
class Client(BaseClient):
|
class Client(BaseClient):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
|
|
@ -97,7 +140,6 @@ class Client(BaseClient):
|
||||||
async def async_images(self) -> Images:
|
async def async_images(self) -> Images:
|
||||||
return self._images
|
return self._images
|
||||||
|
|
||||||
|
|
||||||
class Completions:
|
class Completions:
|
||||||
def __init__(self, client: Client, provider: ProviderType = None):
|
def __init__(self, client: Client, provider: ProviderType = None):
|
||||||
self.client: Client = client
|
self.client: Client = client
|
||||||
|
|
@ -129,25 +171,115 @@ class Completions:
|
||||||
)
|
)
|
||||||
|
|
||||||
stop = [stop] if isinstance(stop, str) else stop
|
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
|
|
||||||
)
|
|
||||||
|
|
||||||
response = iter_response(response, stream, response_format, max_tokens, stop)
|
|
||||||
response = iter_append_model_and_provider(response)
|
|
||||||
|
|
||||||
return response if stream else next(response)
|
|
||||||
|
|
||||||
|
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,
|
||||||
|
self.provider if provider is None else provider,
|
||||||
|
stream,
|
||||||
|
ignored,
|
||||||
|
ignore_working,
|
||||||
|
ignore_stream,
|
||||||
|
)
|
||||||
|
|
||||||
|
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:
|
class Chat:
|
||||||
completions: Completions
|
completions: Completions
|
||||||
|
|
@ -155,153 +287,224 @@ class Chat:
|
||||||
def __init__(self, client: Client, provider: ProviderType = None):
|
def __init__(self, client: Client, provider: ProviderType = None):
|
||||||
self.completions = Completions(client, provider)
|
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
|
||||||
|
|
||||||
def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
|
if not hasattr(response, '__aiter__'):
|
||||||
logging.info("Starting iter_image_response")
|
response = to_async_iterator(response)
|
||||||
response_list = list(response)
|
|
||||||
logging.info(f"Response list: {response_list}")
|
async for chunk in response:
|
||||||
|
if isinstance(chunk, FinishReason):
|
||||||
for chunk in response_list:
|
finish_reason = chunk.reason
|
||||||
logging.info(f"Processing chunk: {chunk}")
|
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:
|
||||||
|
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):
|
if isinstance(chunk, ImageProviderResponse):
|
||||||
logging.info("Found ImageProviderResponse")
|
response_list.extend(chunk.get_list())
|
||||||
return ImagesResponse([Image(image) for image in chunk.get_list()])
|
elif isinstance(chunk, str):
|
||||||
|
response_list.append(chunk)
|
||||||
logging.warning("No ImageProviderResponse found in the response")
|
|
||||||
|
if response_list:
|
||||||
|
return ImagesResponse([Image(image) for image in response_list])
|
||||||
|
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
async def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:
|
||||||
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":
|
if isinstance(provider, type) and provider.__name__ == "You":
|
||||||
kwargs["chat_mode"] = "create"
|
kwargs["chat_mode"] = "create"
|
||||||
else:
|
else:
|
||||||
prompt = f"create an image with: {prompt}"
|
prompt = f"create an image with: {prompt}"
|
||||||
|
|
||||||
response = provider.create_completion(
|
if asyncio.iscoroutinefunction(provider.create_completion):
|
||||||
model,
|
response = await provider.create_completion(
|
||||||
[{"role": "user", "content": prompt}],
|
model,
|
||||||
stream=True,
|
[{"role": "user", "content": prompt}],
|
||||||
proxy=client.get_proxy(),
|
stream=True,
|
||||||
**kwargs
|
proxy=client.get_proxy(),
|
||||||
)
|
**kwargs
|
||||||
|
)
|
||||||
logging.info(f"Response from create_completion: {response}")
|
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
|
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:
|
class Images:
|
||||||
def __init__(self, client: 'Client', provider: ImageProvider = None):
|
def __init__(self, client: 'Client', provider: 'ImageProvider' = None):
|
||||||
self.client: 'Client' = client
|
self.client: 'Client' = client
|
||||||
self.provider: ImageProvider = provider
|
self.provider: 'ImageProvider' = provider
|
||||||
self.models: ImageModels = ImageModels(client)
|
self.models: ImageModels = ImageModels(client)
|
||||||
|
|
||||||
def generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse:
|
def generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
|
||||||
logging.info(f"Starting synchronous image generation for model: {model}, prompt: {prompt}")
|
"""
|
||||||
try:
|
Synchronous generate method that runs the async_generate method in an event loop.
|
||||||
loop = asyncio.get_event_loop()
|
"""
|
||||||
except RuntimeError:
|
return asyncio.run(self.async_generate(prompt, model, response_format=response_format, **kwargs))
|
||||||
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:
|
async def async_generate(self, prompt: str, model: str = None, response_format: str = "url", **kwargs) -> ImagesResponse:
|
||||||
logging.info(f"Generating image for model: {model}, prompt: {prompt}")
|
|
||||||
provider = self.models.get(model, self.provider)
|
provider = self.models.get(model, self.provider)
|
||||||
if provider is None:
|
if provider is None:
|
||||||
raise ValueError(f"Unknown model: {model}")
|
raise ValueError(f"Unknown model: {model}")
|
||||||
|
|
||||||
logging.info(f"Provider: {provider}")
|
|
||||||
|
|
||||||
if isinstance(provider, IterListProvider):
|
if isinstance(provider, IterListProvider):
|
||||||
if provider.providers:
|
if provider.providers:
|
||||||
provider = provider.providers[0]
|
provider = provider.providers[0]
|
||||||
logging.info(f"Using first provider from IterListProvider: {provider}")
|
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"IterListProvider for model {model} has no providers")
|
raise ValueError(f"IterListProvider for model {model} has no providers")
|
||||||
|
|
||||||
if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
|
if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
|
||||||
logging.info("Using AsyncGeneratorProvider")
|
|
||||||
messages = [{"role": "user", "content": prompt}]
|
messages = [{"role": "user", "content": prompt}]
|
||||||
async for response in provider.create_async_generator(model, messages, **kwargs):
|
async for response in provider.create_async_generator(model, messages, **kwargs):
|
||||||
if isinstance(response, ImageResponse):
|
if isinstance(response, ImageResponse):
|
||||||
return self._process_image_response(response)
|
return await self._process_image_response(response, response_format)
|
||||||
elif isinstance(response, str):
|
elif isinstance(response, str):
|
||||||
image_response = ImageResponse([response], prompt)
|
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'):
|
elif hasattr(provider, 'create'):
|
||||||
logging.info("Using provider's create method")
|
|
||||||
if asyncio.iscoroutinefunction(provider.create):
|
if asyncio.iscoroutinefunction(provider.create):
|
||||||
response = await provider.create(prompt)
|
response = await provider.create(prompt)
|
||||||
else:
|
else:
|
||||||
response = provider.create(prompt)
|
response = provider.create(prompt)
|
||||||
|
|
||||||
if isinstance(response, ImageResponse):
|
if isinstance(response, ImageResponse):
|
||||||
return self._process_image_response(response)
|
return await self._process_image_response(response, response_format)
|
||||||
elif isinstance(response, str):
|
elif isinstance(response, str):
|
||||||
image_response = ImageResponse([response], prompt)
|
image_response = ImageResponse([response], prompt)
|
||||||
return self._process_image_response(image_response)
|
return await self._process_image_response(image_response, response_format)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Provider {provider} does not support image generation")
|
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)}")
|
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 = []
|
processed_images = []
|
||||||
|
|
||||||
for image_data in response.get_list():
|
for image_data in response.get_list():
|
||||||
if image_data.startswith('http://') or image_data.startswith('https://'):
|
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:
|
else:
|
||||||
image = to_image(image_data)
|
# Assume image_data is base64 data or binary
|
||||||
file_name = self._save_image(image)
|
if response_format == "url":
|
||||||
processed_images.append(Image(url=file_name))
|
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)
|
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:
|
def _save_image(self, image: 'PILImage') -> str:
|
||||||
os.makedirs('generated_images', exist_ok=True)
|
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)
|
image.save(file_name)
|
||||||
return file_name
|
return file_name
|
||||||
|
|
||||||
async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs):
|
async def create_variation(self, image: Union[str, bytes], model: str = None, response_format: str = "url", **kwargs):
|
||||||
provider = self.models.get(model, self.provider)
|
# Existing implementation, adjust if you want to support b64_json here as well
|
||||||
if provider is None:
|
pass
|
||||||
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")
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue