From b8d12fc78761c25088593172446b4fa2cec51916 Mon Sep 17 00:00:00 2001
From: kqlio67 <kqlio67@users.noreply.github.com>
Date: Thu, 5 Sep 2024 18:10:19 +0300
Subject: feat(AiChats): add DALL-E support and improve error handling

---
 g4f/Provider/AiChats.py | 69 +++++++++++++++++++++++++++++++++----------------
 1 file 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 ""
-- 
cgit v1.2.3