返回提交历史
Modified
g4f/client/client.py
+170
-41
XFEstudio/gpt4free
feat(images): add async support and logging, enhance image handling
44ba42f5
代码差异
1 个文件
+170
-41
@@ -1,9 +1,14 @@
1
1
from __future__ import annotations
2
2
3
import os
3
4
import time
4
5
import random
5
6
import string
6
7
import logging
8
import asyncio
9
from typing import Union
10
from ..providers.base_provider import AsyncGeneratorProvider
11
from ..image import ImageResponse, to_image, to_data_uri
7
12
from ..typing import Union, Iterator, Messages, ImageType
8
13
from ..providers.types import BaseProvider, ProviderType, FinishReason
9
14
from ..providers.conversation import BaseConversation
@@ -15,9 +20,12 @@ from .types import IterResponse, ImageProvider
15
20
from .types import Client as BaseClient
16
21
from .service import get_model_and_provider, get_last_provider
17
22
from .helper import find_stop, filter_json, filter_none
23
from ..models import ModelUtils
24
from ..Provider import IterListProvider
25
18
26
19
27
def iter_response(
20
response: iter[str],
28
response: Iterator[str],
21
29
stream: bool,
22
30
response_format: dict = None,
23
31
max_tokens: int = None,
@@ -26,6 +34,7 @@ def iter_response(
26
34
content = ""
27
35
finish_reason = None
28
36
completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
37
29
38
for idx, chunk in enumerate(response):
30
39
if isinstance(chunk, FinishReason):
31
40
finish_reason = chunk.reason
@@ -33,17 +42,25 @@ def iter_response(
33
42
elif isinstance(chunk, BaseConversation):
34
43
yield chunk
35
44
continue
45
36
46
content += str(chunk)
47
37
48
if max_tokens is not None and idx + 1 >= max_tokens:
38
49
finish_reason = "length"
50
39
51
first, content, chunk = find_stop(stop, content, chunk if stream else None)
52
40
53
if first != -1:
41
54
finish_reason = "stop"
55
42
56
if stream:
43
57
yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
58
44
59
if finish_reason is not None:
45
60
break
61
46
62
finish_reason = "stop" if finish_reason is None else finish_reason
63
47
64
if stream:
48
65
yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
49
66
else:
@@ -52,14 +69,16 @@ def iter_response(
52
69
content = filter_json(content)
53
70
yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
54
71
72
55
73
def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
56
74
last_provider = None
57
75
for chunk in response:
58
76
last_provider = get_last_provider(True) if last_provider is None else last_provider
59
77
chunk.model = last_provider.get("model")
60
chunk.provider = last_provider.get("name")
78
chunk.provider = last_provider.get("name")
61
79
yield chunk
62
80
81
63
82
class Client(BaseClient):
64
83
def __init__(
65
84
self,
@@ -69,9 +88,17 @@ class Client(BaseClient):
69
88
) -> None:
70
89
super().__init__(**kwargs)
71
90
self.chat: Chat = Chat(self, provider)
72
self.images: Images = Images(self, image_provider)
91
self._images: Images = Images(self, image_provider)
73
92
74
class Completions():
93
@property
94
def images(self) -> Images:
95
return self._images
96
97
async def async_images(self) -> Images:
98
return self._images
99
100
101
class Completions:
75
102
def __init__(self, client: Client, provider: ProviderType = None):
76
103
self.client: Client = client
77
104
self.provider: ProviderType = provider
@@ -87,7 +114,7 @@ class Completions():
87
114
max_tokens: int = None,
88
115
stop: Union[list[str], str] = None,
89
116
api_key: str = None,
90
ignored : list[str] = None,
117
ignored: list[str] = None,
91
118
ignore_working: bool = False,
92
119
ignore_stream: bool = False,
93
120
**kwargs
@@ -100,11 +127,13 @@ class Completions():
100
127
ignore_working,
101
128
ignore_stream,
102
129
)
103
130
104
131
stop = [stop] if isinstance(stop, str) else stop
132
105
133
response = provider.create_completion(
106
model, messages,
107
stream=stream,
134
model,
135
messages,
136
stream=stream,
108
137
**filter_none(
109
138
proxy=self.client.get_proxy() if proxy is None else proxy,
110
139
max_tokens=max_tokens,
@@ -113,66 +142,166 @@ class Completions():
113
142
),
114
143
**kwargs
115
144
)
145
116
146
response = iter_response(response, stream, response_format, max_tokens, stop)
117
147
response = iter_append_model_and_provider(response)
148
118
149
return response if stream else next(response)
119
150
120
class Chat():
151
152
class Chat:
121
153
completions: Completions
122
154
123
155
def __init__(self, client: Client, provider: ProviderType = None):
124
156
self.completions = Completions(client, provider)
125
157
158
126
159
def iter_image_response(response: Iterator) -> Union[ImagesResponse, None]:
127
for chunk in list(response):
160
logging.info("Starting iter_image_response")
161
response_list = list(response)
162
logging.info(f"Response list: {response_list}")
163
164
for chunk in response_list:
165
logging.info(f"Processing chunk: {chunk}")
128
166
if isinstance(chunk, ImageProviderResponse):
167
logging.info("Found ImageProviderResponse")
129
168
return ImagesResponse([Image(image) for image in chunk.get_list()])
130
131
def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator:
169
170
logging.warning("No ImageProviderResponse found in the response")
171
return None
132
172
133
173
174
def create_image(client: Client, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> Iterator:
175
logging.info(f"Creating image with provider: {provider}, model: {model}, prompt: {prompt}")
176
134
177
if isinstance(provider, type) and provider.__name__ == "You":
135
178
kwargs["chat_mode"] = "create"
136
179
else:
137
prompt = f"create a image with: {prompt}"
138
return provider.create_completion(
180
prompt = f"create an image with: {prompt}"
181
182
response = provider.create_completion(
139
183
model,
140
184
[{"role": "user", "content": prompt}],
141
185
stream=True,
142
186
proxy=client.get_proxy(),
143
187
**kwargs
144
188
)
189
190
logging.info(f"Response from create_completion: {response}")
191
return response
145
192
146
class Images():
147
def __init__(self, client: Client, provider: ImageProvider = None):
148
self.client: Client = client
193
194
class Images:
195
def __init__(self, client: 'Client', provider: ImageProvider = None):
196
self.client: 'Client' = client
149
197
self.provider: ImageProvider = provider
150
198
self.models: ImageModels = ImageModels(client)
151
199
152
def generate(self, prompt, model: str = None, **kwargs) -> ImagesResponse:
200
def generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse:
201
logging.info(f"Starting synchronous image generation for model: {model}, prompt: {prompt}")
202
try:
203
loop = asyncio.get_event_loop()
204
except RuntimeError:
205
loop = asyncio.new_event_loop()
206
asyncio.set_event_loop(loop)
207
208
try:
209
result = loop.run_until_complete(self.async_generate(prompt, model, **kwargs))
210
logging.info(f"Synchronous image generation completed. Result: {result}")
211
return result
212
except Exception as e:
213
logging.error(f"Error in synchronous image generation: {str(e)}")
214
raise
215
finally:
216
if loop.is_running():
217
loop.close()
218
219
async def async_generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse:
220
logging.info(f"Generating image for model: {model}, prompt: {prompt}")
153
221
provider = self.models.get(model, self.provider)
154
if isinstance(provider, type) and issubclass(provider, BaseProvider):
155
response = create_image(self.client, provider, prompt, **kwargs)
222
if provider is None:
223
raise ValueError(f"Unknown model: {model}")
224
225
logging.info(f"Provider: {provider}")
226
227
if isinstance(provider, IterListProvider):
228
if provider.providers:
229
provider = provider.providers[0]
230
logging.info(f"Using first provider from IterListProvider: {provider}")
231
else:
232
raise ValueError(f"IterListProvider for model {model} has no providers")
233
234
if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
235
logging.info("Using AsyncGeneratorProvider")
236
messages = [{"role": "user", "content": prompt}]
237
async for response in provider.create_async_generator(model, messages, **kwargs):
238
if isinstance(response, ImageResponse):
239
return self._process_image_response(response)
240
elif isinstance(response, str):
241
image_response = ImageResponse([response], prompt)
242
return self._process_image_response(image_response)
243
elif hasattr(provider, 'create'):
244
logging.info("Using provider's create method")
245
if asyncio.iscoroutinefunction(provider.create):
246
response = await provider.create(prompt)
247
else:
248
response = provider.create(prompt)
249
250
if isinstance(response, ImageResponse):
251
return self._process_image_response(response)
252
elif isinstance(response, str):
253
image_response = ImageResponse([response], prompt)
254
return self._process_image_response(image_response)
156
255
else:
157
response = list(provider.create(prompt))
158
image = iter_image_response(response)
159
if image is None:
160
raise NoImageResponseError()
161
return image
256
raise ValueError(f"Provider {provider} does not support image generation")
257
258
logging.error(f"Unexpected response type: {type(response)}")
259
raise NoImageResponseError(f"Unexpected response type: {type(response)}")
260
261
def _process_image_response(self, response: ImageResponse) -> ImagesResponse:
262
processed_images = []
263
for image_data in response.get_list():
264
if image_data.startswith('http://') or image_data.startswith('https://'):
265
processed_images.append(Image(url=image_data))
266
else:
267
image = to_image(image_data)
268
file_name = self._save_image(image)
269
processed_images.append(Image(url=file_name))
270
return ImagesResponse(processed_images)
271
272
def _save_image(self, image: 'PILImage') -> str:
273
os.makedirs('generated_images', exist_ok=True)
274
file_name = f"generated_images/image_{int(time.time())}.png"
275
image.save(file_name)
276
return file_name
162
277
163
def create_variation(self, image: ImageType, model: str = None, **kwargs):
278
async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs):
164
279
provider = self.models.get(model, self.provider)
165
result = None
166
if isinstance(provider, type) and issubclass(provider, BaseProvider):
167
response = provider.create_completion(
168
"",
169
[{"role": "user", "content": "create a image like this"}],
170
True,
171
image=image,
172
proxy=self.client.get_proxy(),
173
**kwargs
174
)
175
result = iter_image_response(response)
176
if result is None:
177
raise NoImageResponseError()
178
return result
280
if provider is None:
281
raise ValueError(f"Unknown model: {model}")
282
283
if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):
284
messages = [{"role": "user", "content": "create a variation of this image"}]
285
image_data = to_data_uri(image)
286
async for response in provider.create_async_generator(model, messages, image=image_data, **kwargs):
287
if isinstance(response, ImageResponse):
288
return self._process_image_response(response)
289
elif isinstance(response, str):
290
image_response = ImageResponse([response], "Image variation")
291
return self._process_image_response(image_response)
292
elif hasattr(provider, 'create_variation'):
293
if asyncio.iscoroutinefunction(provider.create_variation):
294
response = await provider.create_variation(image, **kwargs)
295
else:
296
response = provider.create_variation(image, **kwargs)
297
298
if isinstance(response, ImageResponse):
299
return self._process_image_response(response)
300
elif isinstance(response, str):
301
image_response = ImageResponse([response], "Image variation")
302
return self._process_image_response(image_response)
303
else:
304
raise ValueError(f"Provider {provider} does not support image variation")
305
306
raise NoImageResponseError("Failed to create image variation")
307