feat(g4f/client/async_client.py): enhance image generation and logging

@@ -1,32 +1,37 @@

1 1 from __future__ import annotations

2 2

3 import os

3 4 import time

4 5 import random

5 6 import string

7 import logging

6 8 import asyncio

7 import base64

8 from aiohttp import ClientSession, BaseConnector

9

10 from .types import Client as BaseClient

11 from .types import ProviderType, FinishReason

12 from .stubs import ChatCompletion, ChatCompletionChunk, ImagesResponse, Image

13 from .types import AsyncIterResponse, ImageProvider

14 from .image_models import ImageModels

15 from .helper import filter_json, find_stop, filter_none, cast_iter_async

16 from .service import get_last_provider, get_model_and_provider

17 from ..Provider import ProviderUtils

18 from ..typing import Union, Messages, AsyncIterator, ImageType

19 from ..errors import NoImageResponseError, ProviderNotFoundError

20 from ..requests.aiohttp import get_connector

9 from typing import Union, AsyncIterator

10 from ..providers.base_provider import AsyncGeneratorProvider

11 from ..image import ImageResponse, to_image, to_data_uri

12 from ..typing import Messages, ImageType

13 from ..providers.types import BaseProvider, ProviderType, FinishReason

21 14 from ..providers.conversation import BaseConversation

22 from ..image import ImageResponse as ImageProviderResponse, ImageDataResponse

15 from ..image import ImageResponse as ImageProviderResponse

16 from ..errors import NoImageResponseError

17 from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse

18 from .image_models import ImageModels

19 from .types import IterResponse, ImageProvider

20 from .types import Client as BaseClient

21 from .service import get_model_and_provider, get_last_provider

22 from .helper import find_stop, filter_json, filter_none

23 from ..models import ModelUtils

24 from ..Provider import IterListProvider

25 from .helper import cast_iter_async

23 26

24 27 try:

25 anext

28 anext # Python 3.8+

26 29 except NameError:

27 async def anext(iter):

28 async for chunk in iter:

29 return chunk

30 async def anext(aiter):

31 try:

32 return await aiter.__anext__()

33 except StopAsyncIteration:

34 raise StopIteration

30 35

31 36 async def iter_response(

32 37 response: AsyncIterator[str],

@@ -34,11 +39,12 @@ async def iter_response(

34 39 response_format: dict = None,

35 40 max_tokens: int = None,

36 41 stop: list = None

37 ) -> AsyncIterResponse:

42 ) -> AsyncIterator[Union[ChatCompletion, ChatCompletionChunk]]:

38 43 content = ""

39 44 finish_reason = None

40 45 completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))

41 count: int = 0

46 idx = 0

47

42 48 async for chunk in response:

43 49 if isinstance(chunk, FinishReason):

44 50 finish_reason = chunk.reason

@@ -46,18 +52,26 @@ async def iter_response(

46 52 elif isinstance(chunk, BaseConversation):

47 53 yield chunk

48 54 continue

55

49 56 content += str(chunk)

50 count += 1

51 if max_tokens is not None and count >= max_tokens:

57 idx += 1

58

59 if max_tokens is not None and idx >= max_tokens:

52 60 finish_reason = "length"

53 first, content, chunk = find_stop(stop, content, chunk)

61

62 first, content, chunk = find_stop(stop, content, chunk if stream else None)

63

54 64 if first != -1:

55 65 finish_reason = "stop"

66

56 67 if stream:

57 68 yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))

69

58 70 if finish_reason is not None:

59 71 break

72

60 73 finish_reason = "stop" if finish_reason is None else finish_reason

74

61 75 if stream:

62 76 yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))

63 77 else:

@@ -66,12 +80,12 @@ async def iter_response(

66 80 content = filter_json(content)

67 81 yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))

68 82

69 async def iter_append_model_and_provider(response: AsyncIterResponse) -> AsyncIterResponse:

83 async def iter_append_model_and_provider(response: AsyncIterator) -> AsyncIterator:

70 84 last_provider = None

71 85 async for chunk in response:

72 86 last_provider = get_last_provider(True) if last_provider is None else last_provider

73 87 chunk.model = last_provider.get("model")

74 chunk.provider = last_provider.get("name")

88 chunk.provider = last_provider.get("name")

75 89 yield chunk

76 90

77 91 class AsyncClient(BaseClient):

@@ -80,59 +94,32 @@ class AsyncClient(BaseClient):

80 94 provider: ProviderType = None,

81 95 image_provider: ImageProvider = None,

82 96 **kwargs

83 ):

97 ) -> None:

84 98 super().__init__(**kwargs)

85 99 self.chat: Chat = Chat(self, provider)

86 self.images: Images = Images(self, image_provider)

87

88 def create_response(

89 messages: Messages,

90 model: str,

91 provider: ProviderType = None,

92 stream: bool = False,

93 proxy: str = None,

94 max_tokens: int = None,

95 stop: list[str] = None,

96 api_key: str = None,

97 **kwargs

98 ):

99 has_asnyc = hasattr(provider, "create_async_generator")

100 if has_asnyc:

101 create = provider.create_async_generator

102 else:

103 create = provider.create_completion

104 response = create(

105 model, messages,

106 stream=stream,

107 **filter_none(

108 proxy=proxy,

109 max_tokens=max_tokens,

110 stop=stop,

111 api_key=api_key

112 ),

113 **kwargs

114 )

115 if not has_asnyc:

116 response = cast_iter_async(response)

117 return response

100 self._images: Images = Images(self, image_provider)

118 101

119 class Completions():

120 def __init__(self, client: AsyncClient, provider: ProviderType = None):

121 self.client: AsyncClient = client

102 @property

103 def images(self) -> Images:

104 return self._images

105

106 class Completions:

107 def __init__(self, client: 'AsyncClient', provider: ProviderType = None):

108 self.client: 'AsyncClient' = client

122 109 self.provider: ProviderType = provider

123 110

124 def create(

111 async def create(

125 112 self,

126 113 messages: Messages,

127 114 model: str,

128 115 provider: ProviderType = None,

129 116 stream: bool = False,

130 117 proxy: str = None,

118 response_format: dict = None,

131 119 max_tokens: int = None,

132 120 stop: Union[list[str], str] = None,

133 121 api_key: str = None,

134 response_format: dict = None,

135 ignored : list[str] = None,

122 ignored: list[str] = None,

136 123 ignore_working: bool = False,

137 124 ignore_stream: bool = False,

138 125 **kwargs

@@ -143,133 +130,171 @@ class Completions():

143 130 stream,

144 131 ignored,

145 132 ignore_working,

146 ignore_stream

133 ignore_stream,

147 134 )

135

148 136 stop = [stop] if isinstance(stop, str) else stop

149 response = create_response(

150 messages, model,

151 provider, stream,

152 proxy=self.client.get_proxy() if proxy is None else proxy,

153 max_tokens=max_tokens,

154 stop=stop,

155 api_key=self.client.api_key if api_key is None else api_key,

137

138 response = provider.create_completion(

139 model,

140 messages,

141 stream=stream,

142 **filter_none(

143 proxy=self.client.get_proxy() if proxy is None else proxy,

144 max_tokens=max_tokens,

145 stop=stop,

146 api_key=self.client.api_key if api_key is None else api_key

147 ),

156 148 **kwargs

157 149 )

158 response = iter_response(response, stream, response_format, max_tokens, stop)

159 response = iter_append_model_and_provider(response)

160 return response if stream else anext(response)

161 150

162 class Chat():

151 if isinstance(response, AsyncIterator):

152 response = iter_response(response, stream, response_format, max_tokens, stop)

153 response = iter_append_model_and_provider(response)

154 return response if stream else await anext(response)

155 else:

156 response = cast_iter_async(response)

157 response = iter_response(response, stream, response_format, max_tokens, stop)

158 response = iter_append_model_and_provider(response)

159 return response if stream else await anext(response)

160

161

162

163 class Chat:

163 164 completions: Completions

164 165

165 166 def __init__(self, client: AsyncClient, provider: ProviderType = None):

166 167 self.completions = Completions(client, provider)

167 168

168 async def iter_image_response(

169 response: AsyncIterator,

170 response_format: str = None,

171 connector: BaseConnector = None,

172 proxy: str = None

173 ) -> Union[ImagesResponse, None]:

169 async def iter_image_response(response: AsyncIterator) -> Union[ImagesResponse, None]:

170 logging.info("Starting iter_image_response")

174 171 async for chunk in response:

172 logging.info(f"Processing chunk: {chunk}")

175 173 if isinstance(chunk, ImageProviderResponse):

176 if response_format == "b64_json":

177 async with ClientSession(

178 connector=get_connector(connector, proxy),

179 cookies=chunk.options.get("cookies")

180 ) as session:

181 async def fetch_image(image):

182 async with session.get(image) as response:

183 return base64.b64encode(await response.content.read()).decode()

184 images = await asyncio.gather(*[fetch_image(image) for image in chunk.get_list()])

185 return ImagesResponse([Image(None, image, chunk.alt) for image in images], int(time.time()))

186 return ImagesResponse([Image(image, None, chunk.alt) for image in chunk.get_list()], int(time.time()))

187 elif isinstance(chunk, ImageDataResponse):

188 return ImagesResponse([Image(None, image, chunk.alt) for image in chunk.get_list()], int(time.time()))

189

190 def create_image(provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:

174 logging.info("Found ImageProviderResponse")

175 return ImagesResponse([Image(image) for image in chunk.get_list()])

176

177 logging.warning("No ImageProviderResponse found in the response")

178 return None

179

180

async def create_image(client: AsyncClient, provider: ProviderType, prompt: str, model: str = "", **kwargs) -> AsyncIterator:

181 logging.info(f"Creating image with provider: {provider}, model: {model}, prompt: {prompt}")

182

191 183 if isinstance(provider, type) and provider.__name__ == "You":

192 184 kwargs["chat_mode"] = "create"

193 185 else:

194 prompt = f"create a image with: {prompt}"

195 return provider.create_async_generator(

186 prompt = f"create an image with: {prompt}"

187

188 response = await provider.create_completion(

196 189 model,

197 190 [{"role": "user", "content": prompt}],

198 191 stream=True,

192 proxy=client.get_proxy(),

199 193 **kwargs

200 194 )

195

196 logging.info(f"Response from create_completion: {response}")

197 return response

201 198

202 class Images():

203 def __init__(self, client: AsyncClient, provider: ImageProvider = None):

204 self.client: AsyncClient = client

199 class Images:

200 def __init__(self, client: 'AsyncClient', provider: ImageProvider = None):

201 self.client: 'AsyncClient' = client

205 202 self.provider: ImageProvider = provider

206 203 self.models: ImageModels = ImageModels(client)

207 204

208 def get_provider(self, model: str, provider: ProviderType = None):

209 if isinstance(provider, str):

210 if provider in ProviderUtils.convert:

211 provider = ProviderUtils.convert[provider]

205 async def generate(self, prompt: str, model: str = None, **kwargs) -> ImagesResponse:

206 logging.info(f"Starting asynchronous image generation for model: {model}, prompt: {prompt}")

207 provider = self.models.get(model, self.provider)

208 if provider is None:

209 raise ValueError(f"Unknown model: {model}")

210

211 logging.info(f"Provider: {provider}")

212

213 if isinstance(provider, IterListProvider):

214 if provider.providers:

215 provider = provider.providers[0]

216 logging.info(f"Using first provider from IterListProvider: {provider}")

217 else:

218 raise ValueError(f"IterListProvider for model {model} has no providers")

219

220 if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):

221 logging.info("Using AsyncGeneratorProvider")

222 messages = [{"role": "user", "content": prompt}]

223 async for response in provider.create_async_generator(model, messages, **kwargs):

224 if isinstance(response, ImageResponse):

225 return self._process_image_response(response)

226 elif isinstance(response, str):

227 image_response = ImageResponse([response], prompt)

228 return self._process_image_response(image_response)

229 elif hasattr(provider, 'create'):

230 logging.info("Using provider's create method")

231 async_create = asyncio.iscoroutinefunction(provider.create)

232 if async_create:

233 response = await provider.create(prompt)

212 234 else:

213 raise ProviderNotFoundError(f'Provider not found: {provider}')

235 response = provider.create(prompt)

236

237 if isinstance(response, ImageResponse):

238 return self._process_image_response(response)

239 elif isinstance(response, str):

240 image_response = ImageResponse([response], prompt)

241 return self._process_image_response(image_response)

242 elif hasattr(provider, 'create_completion'):

243 logging.info("Using provider's create_completion method")

244 response = await create_image(provider, prompt, model, **kwargs)

245 async for chunk in response:

246 if isinstance(chunk, ImageProviderResponse):

247 logging.info("Found ImageProviderResponse")

248 return ImagesResponse([Image(image) for image in chunk.get_list()])

214 249 else:

215 provider = self.models.get(model, self.provider)

216 return provider

250 raise ValueError(f"Provider {provider} does not support image generation")

217 251

218 async def generate(

219 self,

220 prompt,

221 model: str = "",

222 provider: ProviderType = None,

223 response_format: str = None,

224 connector: BaseConnector = None,

225 proxy: str = None,

226 **kwargs

227 ) -> ImagesResponse:

228 provider = self.get_provider(model, provider)

229 if hasattr(provider, "create_async_generator"):

230 response = create_image(

231 provider,

232 prompt,

233 **filter_none(

234 response_format=response_format,

235 connector=connector,

236 proxy=self.client.get_proxy() if proxy is None else proxy,

237 ),

238 **kwargs

239 )

252 logging.error(f"Unexpected response type: {type(response)}")

253 raise NoImageResponseError(f"Unexpected response type: {type(response)}")

254

255 def _process_image_response(self, response: ImageResponse) -> ImagesResponse:

256 processed_images = []

257 for image_data in response.get_list():

258 if image_data.startswith('http://') or image_data.startswith('https://'):

259 processed_images.append(Image(url=image_data))

260 else:

261 image = to_image(image_data)

262 file_name = self._save_image(image)

263 processed_images.append(Image(url=file_name))

264 return ImagesResponse(processed_images)

265

266 def _save_image(self, image: 'PILImage') -> str:

267 os.makedirs('generated_images', exist_ok=True)

268 file_name = f"generated_images/image_{int(time.time())}.png"

269 image.save(file_name)

270 return file_name

271

272 async def create_variation(self, image: Union[str, bytes], model: str = None, **kwargs) -> ImagesResponse:

273 provider = self.models.get(model, self.provider)

274 if provider is None:

275 raise ValueError(f"Unknown model: {model}")

276

277 if isinstance(provider, type) and issubclass(provider, AsyncGeneratorProvider):

278 messages = [{"role": "user", "content": "create a variation of this image"}]

279 image_data = to_data_uri(image)

280 async for response in provider.create_async_generator(model, messages, image=image_data, **kwargs):

281 if isinstance(response, ImageResponse):

282 return self._process_image_response(response)

283 elif isinstance(response, str):

284 image_response = ImageResponse([response], "Image variation")

285 return self._process_image_response(image_response)

286 elif hasattr(provider, 'create_variation'):

287 if asyncio.iscoroutinefunction(provider.create_variation):

288 response = await provider.create_variation(image, **kwargs)

289 else:

290 response = provider.create_variation(image, **kwargs)

291

292 if isinstance(response, ImageResponse):

293 return self._process_image_response(response)

294 elif isinstance(response, str):

295 image_response = ImageResponse([response], "Image variation")

296 return self._process_image_response(image_response)

240 297 else:

241 response = await provider.create_async(prompt)

242 return ImagesResponse([Image(image) for image in response.get_list()])

243 image = await iter_image_response(response, response_format, connector, proxy)

244 if image is None:

245 raise NoImageResponseError()

246 return image

247

248 async def create_variation(

249 self,

250 image: ImageType,

251 model: str = None,

252 response_format: str = None,

253 connector: BaseConnector = None,

254 proxy: str = None,

255 **kwargs

256 ):

257 provider = self.get_provider(model, provider)

258 result = None

259 if hasattr(provider, "create_async_generator"):

260 response = provider.create_async_generator(

261 "",

262 [{"role": "user", "content": "create a image like this"}],

263 stream=True,

264 image=image,

265 **filter_none(

266 response_format=response_format,

267 connector=connector,

268 proxy=self.client.get_proxy() if proxy is None else proxy,

269 ),

270 **kwargs

271 )

272 result = iter_image_response(response, response_format, connector, proxy)

273 if result is None:

274 raise NoImageResponseError()

275 return result

298 raise ValueError(f"Provider {provider} does not support image variation")

299

300 raise NoImageResponseError("Failed to create image variation")

gpt4free

feat(g4f/client/async_client.py): enhance image generation and logging

代码差异