diff options
Diffstat (limited to 'g4f')
| -rw-r--r-- | g4f/Provider/AiChats.py | 69 |
1 files changed, 47 insertions, 22 deletions
diff --git a/g4f/Provider/AiChats.py b/g4f/Provider/AiChats.py index f3eb3fc0..8dc7cfcf 100644 --- a/g4f/Provider/AiChats.py +++ b/g4f/Provider/AiChats.py @@ -1,19 +1,20 @@ from __future__ import annotations -from aiohttp import ClientSession import json - +import base64 +from aiohttp import ClientSession from ..typing import AsyncResult, Messages -from .base_provider import AsyncGeneratorProvider -from .helper import format_prompt - +from .base_provider import AsyncGeneratorProvider, ProviderModelMixin +from ..image import ImageResponse -class AiChats(AsyncGeneratorProvider): +class AiChats(AsyncGeneratorProvider, ProviderModelMixin): url = "https://ai-chats.org" api_endpoint = "https://ai-chats.org/chat/send2/" working = True supports_gpt_4 = True supports_message_history = True + default_model = 'gpt-4' + models = ['gpt-4', 'dalle'] @classmethod async def create_async_generator( @@ -23,7 +24,6 @@ class AiChats(AsyncGeneratorProvider): proxy: str = None, **kwargs ) -> AsyncResult: - headers = { "accept": "application/json, text/event-stream", "accept-language": "en-US,en;q=0.9", @@ -31,7 +31,7 @@ class AiChats(AsyncGeneratorProvider): "content-type": "application/json", "origin": cls.url, "pragma": "no-cache", - "referer": f"{cls.url}/chat/", + "referer": f"{cls.url}/{'image' if model == 'dalle' else 'chat'}/", "sec-ch-ua": '"Chromium";v="127", "Not)A;Brand";v="99"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"Linux"', @@ -42,9 +42,9 @@ class AiChats(AsyncGeneratorProvider): } async with ClientSession(headers=headers) as session: - prompt = format_prompt(messages) + prompt = cls.format_prompt(messages) data = { - "type": "chat", + "type": "image" if model == 'dalle' else "chat", "messagesHistory": [ { "from": "you", @@ -52,15 +52,40 @@ class AiChats(AsyncGeneratorProvider): } ] } - async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response: - response.raise_for_status() - full_response = await response.text() - - message = "" - for line in full_response.split('\n'): - if line.startswith('data: ') and line != 'data: ': - message += line[6:] - - message = message.strip() - - yield message + + 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) + else: + yield f"Error: No image URL found in the response. Full response: {response_json}" + else: + yield f"Error: Unexpected response format. Full response: {response_json}" + else: + full_response = await response.text() + message = "" + 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 "" |