1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
from __future__ import annotations
import re
import asyncio
from .. import debug
from ..typing import CreateResult, Messages
from ..base_provider import BaseProvider, ProviderType
system_message = """
You can generate custom images with the DALL-E 3 image generator.
To generate an image with a prompt, do this:
<img data-prompt=\"keywords for the image\">
Don't use images with data uri. It is important to use a prompt instead.
<img data-prompt=\"image caption\">
"""
class CreateImagesProvider(BaseProvider):
"""
Provider class for creating images based on text prompts.
This provider handles image creation requests embedded within message content,
using provided image creation functions.
Attributes:
provider (ProviderType): The underlying provider to handle non-image related tasks.
create_images (callable): A function to create images synchronously.
create_images_async (callable): A function to create images asynchronously.
system_message (str): A message that explains the image creation capability.
include_placeholder (bool): Flag to determine whether to include the image placeholder in the output.
__name__ (str): Name of the provider.
url (str): URL of the provider.
working (bool): Indicates if the provider is operational.
supports_stream (bool): Indicates if the provider supports streaming.
"""
def __init__(
self,
provider: ProviderType,
create_images: callable,
create_async: callable,
system_message: str = system_message,
include_placeholder: bool = True
) -> None:
"""
Initializes the CreateImagesProvider.
Args:
provider (ProviderType): The underlying provider.
create_images (callable): Function to create images synchronously.
create_async (callable): Function to create images asynchronously.
system_message (str, optional): System message to be prefixed to messages. Defaults to a predefined message.
include_placeholder (bool, optional): Whether to include image placeholders in the output. Defaults to True.
"""
self.provider = provider
self.create_images = create_images
self.create_images_async = create_async
self.system_message = system_message
self.include_placeholder = include_placeholder
self.__name__ = provider.__name__
self.url = provider.url
self.working = provider.working
self.supports_stream = provider.supports_stream
def create_completion(
self,
model: str,
messages: Messages,
stream: bool = False,
**kwargs
) -> CreateResult:
"""
Creates a completion result, processing any image creation prompts found within the messages.
Args:
model (str): The model to use for creation.
messages (Messages): The messages to process, which may contain image prompts.
stream (bool, optional): Indicates whether to stream the results. Defaults to False.
**kwargs: Additional keywordarguments for the provider.
Yields:
CreateResult: Yields chunks of the processed messages, including image data if applicable.
Note:
This method processes messages to detect image creation prompts. When such a prompt is found,
it calls the synchronous image creation function and includes the resulting image in the output.
"""
messages.insert(0, {"role": "system", "content": self.system_message})
buffer = ""
for chunk in self.provider.create_completion(model, messages, stream, **kwargs):
if isinstance(chunk, str) and buffer or "<" in chunk:
buffer += chunk
if ">" in buffer:
match = re.search(r'<img data-prompt="(.*?)">', buffer)
if match:
placeholder, prompt = match.group(0), match.group(1)
start, append = buffer.split(placeholder, 1)
if start:
yield start
if self.include_placeholder:
yield placeholder
if debug.logging:
print(f"Create images with prompt: {prompt}")
yield from self.create_images(prompt)
if append:
yield append
else:
yield buffer
buffer = ""
else:
yield chunk
async def create_async(
self,
model: str,
messages: Messages,
**kwargs
) -> str:
"""
Asynchronously creates a response, processing any image creation prompts found within the messages.
Args:
model (str): The model to use for creation.
messages (Messages): The messages to process, which may contain image prompts.
**kwargs: Additional keyword arguments for the provider.
Returns:
str: The processed response string, including asynchronously generated image data if applicable.
Note:
This method processes messages to detect image creation prompts. When such a prompt is found,
it calls the asynchronous image creation function and includes the resulting image in the output.
"""
messages.insert(0, {"role": "system", "content": self.system_message})
response = await self.provider.create_async(model, messages, **kwargs)
matches = re.findall(r'(<img data-prompt="(.*?)">)', response)
results = []
placeholders = []
for placeholder, prompt in matches:
if placeholder not in placeholders:
if debug.logging:
print(f"Create images with prompt: {prompt}")
results.append(self.create_images_async(prompt))
placeholders.append(placeholder)
results = await asyncio.gather(*results)
for idx, result in enumerate(results):
placeholder = placeholder[idx]
if self.include_placeholder:
result = placeholder + result
response = response.replace(placeholder, result)
return response
|
