diff options
| -rw-r--r-- | g4f/Provider/AiChats.py | 33 |
1 files changed, 24 insertions, 9 deletions
diff --git a/g4f/Provider/AiChats.py b/g4f/Provider/AiChats.py index 8dc7cfcf..10127d4f 100644 --- a/g4f/Provider/AiChats.py +++ b/g4f/Provider/AiChats.py @@ -6,6 +6,7 @@ from aiohttp import ClientSession from ..typing import AsyncResult, Messages from .base_provider import AsyncGeneratorProvider, ProviderModelMixin from ..image import ImageResponse +from .helper import format_prompt class AiChats(AsyncGeneratorProvider, ProviderModelMixin): url = "https://ai-chats.org" @@ -39,10 +40,15 @@ class AiChats(AsyncGeneratorProvider, ProviderModelMixin): "sec-fetch-mode": "cors", "sec-fetch-site": "same-origin", "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36", + 'cookie': 'muVyak=LSFNvUWqdgKkGprbDBsfieIoEMzjOQ; LSFNvUWqdgKkGprbDBsfieIoEMzjOQ=ac28831b98143847e83dbe004404e619-1725548624-1725548621; muVyak_hits=9; ai-chat-front=9d714d5dc46a6b47607c9a55e7d12a95; _csrf-front=76c23dc0a013e5d1e21baad2e6ba2b5fdab8d3d8a1d1281aa292353f8147b057a%3A2%3A%7Bi%3A0%3Bs%3A11%3A%22_csrf-front%22%3Bi%3A1%3Bs%3A32%3A%22K9lz0ezsNPMNnfpd_8gT5yEeh-55-cch%22%3B%7D', } - + async with ClientSession(headers=headers) as session: - prompt = cls.format_prompt(messages) + if model == 'dalle': + prompt = messages[-1]['content'] if messages else "" + else: + prompt = format_prompt(messages) + data = { "type": "image" if model == 'dalle' else "chat", "messagesHistory": [ @@ -52,21 +58,21 @@ class AiChats(AsyncGeneratorProvider, ProviderModelMixin): } ] } - + try: async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response: response.raise_for_status() - + if model == 'dalle': response_json = await response.json() - + if 'data' in response_json and response_json['data']: image_url = response_json['data'][0].get('url') if image_url: async with session.get(image_url) as img_response: img_response.raise_for_status() image_data = await img_response.read() - + base64_image = base64.b64encode(image_data).decode('utf-8') base64_url = f"data:image/png;base64,{base64_image}" yield ImageResponse(base64_url, prompt) @@ -80,12 +86,21 @@ class AiChats(AsyncGeneratorProvider, ProviderModelMixin): for line in full_response.split('\n'): if line.startswith('data: ') and line != 'data: ': message += line[6:] - + message = message.strip() yield message except Exception as e: yield f"Error occurred: {str(e)}" @classmethod - def format_prompt(cls, messages: Messages) -> str: - return messages[-1]['content'] if messages else "" + async def create_async( + cls, + model: str, + messages: Messages, + proxy: str = None, + **kwargs + ) -> str: + async for response in cls.create_async_generator(model, messages, proxy, **kwargs): + if isinstance(response, ImageResponse): + return response.images[0] + return response |