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
|
from __future__ import annotations
import re
import asyncio
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 a 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):
def __init__(
self,
provider: ProviderType,
create_images: callable,
create_async: callable,
system_message: str = system_message,
include_placeholder: bool = True
) -> None:
self.provider = provider
self.create_images = create_images
self.create_images_async = create_async
self.system_message = system_message
self.__name__ = provider.__name__
self.working = provider.working
self.supports_stream = provider.supports_stream
self.include_placeholder = include_placeholder
if hasattr(provider, "url"):
self.url = provider.url
def create_completion(
self,
model: str,
messages: Messages,
stream: bool = False,
**kwargs
) -> CreateResult:
messages.insert(0, {"role": "system", "content": self.system_message})
buffer = ""
for chunk in self.provider.create_completion(model, messages, stream, **kwargs):
if 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
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:
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:
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
|
