2024-02-12 05:41:27 -05:00
|
|
|
from __future__ import annotations
|
|
|
|
|
2024-02-22 20:35:13 -05:00
|
|
|
import time
|
|
|
|
import random
|
|
|
|
import string
|
2024-02-12 05:41:27 -05:00
|
|
|
|
2024-04-07 04:36:13 -04:00
|
|
|
from ..typing import Union, Iterator, Messages, ImageType
|
|
|
|
from ..providers.types import BaseProvider, ProviderType, FinishReason
|
|
|
|
from ..image import ImageResponse as ImageProviderResponse
|
|
|
|
from ..errors import NoImageResponseError
|
2024-02-14 03:21:57 -05:00
|
|
|
from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse
|
2024-04-07 04:36:13 -04:00
|
|
|
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
|
2024-02-12 05:41:27 -05:00
|
|
|
|
|
|
|
def iter_response(
|
2024-02-22 20:35:13 -05:00
|
|
|
response: iter[str],
|
2024-02-12 05:41:27 -05:00
|
|
|
stream: bool,
|
|
|
|
response_format: dict = None,
|
|
|
|
max_tokens: int = None,
|
|
|
|
stop: list = None
|
2024-02-22 20:35:13 -05:00
|
|
|
) -> IterResponse:
|
2024-02-12 05:41:27 -05:00
|
|
|
content = ""
|
2024-02-14 03:21:57 -05:00
|
|
|
finish_reason = None
|
2024-02-22 20:35:13 -05:00
|
|
|
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
|
2024-02-12 05:41:27 -05:00
|
|
|
for idx, chunk in enumerate(response):
|
2024-04-05 19:05:00 -04:00
|
|
|
if isinstance(chunk, FinishReason):
|
|
|
|
finish_reason = chunk.reason
|
|
|
|
break
|
2024-02-12 05:41:27 -05:00
|
|
|
content += str(chunk)
|
2024-02-14 03:21:57 -05:00
|
|
|
if max_tokens is not None and idx + 1 >= max_tokens:
|
2024-02-21 18:16:58 -05:00
|
|
|
finish_reason = "length"
|
2024-04-07 04:36:13 -04:00
|
|
|
first, content, chunk = find_stop(stop, content, chunk if stream else None)
|
2024-02-12 05:41:27 -05:00
|
|
|
if first != -1:
|
2024-02-14 03:21:57 -05:00
|
|
|
finish_reason = "stop"
|
|
|
|
if stream:
|
2024-02-22 20:35:13 -05:00
|
|
|
yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
|
2024-02-14 03:21:57 -05:00
|
|
|
if finish_reason is not None:
|
2024-02-12 05:41:27 -05:00
|
|
|
break
|
2024-02-22 20:35:13 -05:00
|
|
|
finish_reason = "stop" if finish_reason is None else finish_reason
|
|
|
|
if stream:
|
|
|
|
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
|
|
|
|
else:
|
2024-02-12 05:41:27 -05:00
|
|
|
if response_format is not None and "type" in response_format:
|
|
|
|
if response_format["type"] == "json_object":
|
2024-04-07 04:36:13 -04:00
|
|
|
content = filter_json(content)
|
2024-02-22 20:35:13 -05:00
|
|
|
yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
|
|
|
|
|
|
|
|
def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
|
|
|
|
last_provider = None
|
|
|
|
for chunk in response:
|
|
|
|
last_provider = get_last_provider(True) if last_provider is None else last_provider
|
|
|
|
chunk.model = last_provider.get("model")
|
|
|
|
chunk.provider = last_provider.get("name")
|
|
|
|
yield chunk
|
2024-02-12 05:41:27 -05:00
|
|
|
|
2024-04-07 04:36:13 -04:00
|
|
|
class Client(BaseClient):
|
2024-02-12 05:41:27 -05:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
provider: ProviderType = None,
|
|
|
|
image_provider: ImageProvider = None,
|
|
|
|
**kwargs
|
|
|
|
) -> None:
|
2024-04-07 04:36:13 -04:00
|
|
|
super().__init__(**kwargs)
|
2024-02-23 05:33:38 -05:00
|
|
|
self.chat: Chat = Chat(self, provider)
|
|
|
|
self.images: Images = Images(self, image_provider)
|
2024-02-12 05:41:27 -05:00
|
|
|
|
|
|
|
class Completions():
|
|
|
|
def __init__(self, client: Client, provider: ProviderType = None):
|
|
|
|
self.client: Client = client
|
|
|
|
self.provider: ProviderType = provider
|
|
|
|
|
|
|
|
def create(
|
|
|
|
self,
|
|
|
|
messages: Messages,
|
|
|
|
model: str,
|
|
|
|
provider: ProviderType = None,
|
|
|
|
stream: bool = False,
|
2024-04-07 04:36:13 -04:00
|
|
|
proxy: str = None,
|
2024-02-12 05:41:27 -05:00
|
|
|
response_format: dict = None,
|
|
|
|
max_tokens: int = None,
|
2024-02-22 20:51:10 -05:00
|
|
|
stop: Union[list[str], str] = None,
|
2024-02-23 05:33:38 -05:00
|
|
|
api_key: str = None,
|
2024-03-11 21:06:06 -04:00
|
|
|
ignored : list[str] = None,
|
|
|
|
ignore_working: bool = False,
|
|
|
|
ignore_stream: bool = False,
|
2024-02-12 05:41:27 -05:00
|
|
|
**kwargs
|
2024-02-27 05:55:40 -05:00
|
|
|
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
|
2024-02-12 05:41:27 -05:00
|
|
|
model, provider = get_model_and_provider(
|
|
|
|
model,
|
|
|
|
self.provider if provider is None else provider,
|
|
|
|
stream,
|
2024-03-11 21:06:06 -04:00
|
|
|
ignored,
|
|
|
|
ignore_working,
|
|
|
|
ignore_stream,
|
2024-02-12 05:41:27 -05:00
|
|
|
)
|
2024-04-12 12:06:54 -04:00
|
|
|
|
2024-02-14 03:21:57 -05:00
|
|
|
stop = [stop] if isinstance(stop, str) else stop
|
2024-02-23 05:33:38 -05:00
|
|
|
response = provider.create_completion(
|
2024-04-07 10:02:49 -04:00
|
|
|
model, messages,
|
|
|
|
stream=stream,
|
2024-02-27 05:55:40 -05:00
|
|
|
**filter_none(
|
2024-04-07 04:36:13 -04:00
|
|
|
proxy=self.client.get_proxy() if proxy is None else proxy,
|
2024-02-27 05:55:40 -05:00
|
|
|
max_tokens=max_tokens,
|
|
|
|
stop=stop,
|
|
|
|
api_key=self.client.api_key if api_key is None else api_key
|
|
|
|
),
|
2024-02-23 05:33:38 -05:00
|
|
|
**kwargs
|
|
|
|
)
|
|
|
|
response = iter_response(response, stream, response_format, max_tokens, stop)
|
|
|
|
response = iter_append_model_and_provider(response)
|
2024-02-12 05:41:27 -05:00
|
|
|
return response if stream else next(response)
|
|
|
|
|
|
|
|
class Chat():
|
|
|
|
completions: Completions
|
|
|
|
|
|
|
|
def __init__(self, client: Client, provider: ProviderType = None):
|
|
|
|
self.completions = Completions(client, provider)
|
2024-02-14 03:21:57 -05:00
|
|
|
|
2024-03-10 21:41:59 -04:00
|
|
|
def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
|
|
|
|
for chunk in list(response):
|
|
|
|
if isinstance(chunk, ImageProviderResponse):
|
|
|
|
return ImagesResponse([Image(image) for image in chunk.get_list()])
|
|
|
|
|
|
|
|
def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator:
|
|
|
|
prompt = f"create a image with: {prompt}"
|
2024-04-07 04:36:13 -04:00
|
|
|
if provider.__name__ == "You":
|
|
|
|
kwargs["chat_mode"] = "create"
|
2024-03-10 21:41:59 -04:00
|
|
|
return provider.create_completion(
|
|
|
|
model,
|
|
|
|
[{"role": "user", "content": prompt}],
|
2024-04-07 04:36:13 -04:00
|
|
|
stream=True,
|
2024-03-10 21:41:59 -04:00
|
|
|
proxy=client.get_proxy(),
|
|
|
|
**kwargs
|
|
|
|
)
|
|
|
|
|
2024-02-12 05:41:27 -05:00
|
|
|
class Images():
|
|
|
|
def __init__(self, client: Client, provider: ImageProvider = None):
|
|
|
|
self.client: Client = client
|
|
|
|
self.provider: ImageProvider = provider
|
|
|
|
self.models: ImageModels = ImageModels(client)
|
|
|
|
|
2024-03-10 21:41:59 -04:00
|
|
|
def generate(self, prompt, model: str = None, **kwargs) -> ImagesResponse:
|
2024-02-14 03:21:57 -05:00
|
|
|
provider = self.models.get(model, self.provider)
|
2024-03-10 21:41:59 -04:00
|
|
|
if isinstance(provider, type) and issubclass(provider, BaseProvider):
|
|
|
|
response = create_image(self.client, provider, prompt, **kwargs)
|
2024-02-12 05:41:27 -05:00
|
|
|
else:
|
2024-04-07 04:36:13 -04:00
|
|
|
response = list(provider.create(prompt))
|
2024-03-10 21:41:59 -04:00
|
|
|
image = iter_image_response(response)
|
|
|
|
if image is None:
|
|
|
|
raise NoImageResponseError()
|
|
|
|
return image
|
2024-02-12 05:41:27 -05:00
|
|
|
|
|
|
|
def create_variation(self, image: ImageType, model: str = None, **kwargs):
|
2024-02-14 03:21:57 -05:00
|
|
|
provider = self.models.get(model, self.provider)
|
2024-02-12 06:08:08 -05:00
|
|
|
result = None
|
|
|
|
if isinstance(provider, type) and issubclass(provider, BaseProvider):
|
2024-02-12 05:41:27 -05:00
|
|
|
response = provider.create_completion(
|
|
|
|
"",
|
|
|
|
[{"role": "user", "content": "create a image like this"}],
|
|
|
|
True,
|
|
|
|
image=image,
|
2024-02-12 06:08:08 -05:00
|
|
|
proxy=self.client.get_proxy(),
|
|
|
|
**kwargs
|
2024-02-12 05:41:27 -05:00
|
|
|
)
|
2024-04-07 04:36:13 -04:00
|
|
|
result = iter_image_response(response)
|
2024-02-12 06:08:08 -05:00
|
|
|
if result is None:
|
|
|
|
raise NoImageResponseError()
|
|
|
|
return result
|