1 files changed, 47 insertions, 109 deletions

diff --git a/g4f/Provider/Nexra.py b/g4f/Provider/Nexra.py
index e2c3e197..65c50e73 100644
--- a/g4f/Provider/Nexra.py
+++ b/g4f/Provider/Nexra.py
@@ -1,40 +1,32 @@
 from __future__ import annotations
-
 import json
-import base64
 from aiohttp import ClientSession
-from typing import AsyncGenerator
 
 from ..typing import AsyncResult, Messages
 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
-from ..image import ImageResponse
 from .helper import format_prompt
+from ..image import ImageResponse
 
 class Nexra(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://nexra.aryahcr.cc"
-    api_endpoint_text = "https://nexra.aryahcr.cc/api/chat/gpt"
-    api_endpoint_image = "https://nexra.aryahcr.cc/api/image/complements"
+    chat_api_endpoint = "https://nexra.aryahcr.cc/api/chat/gpt"
+    image_api_endpoint = "https://nexra.aryahcr.cc/api/image/complements"
     working = True
     supports_gpt_35_turbo = True
     supports_gpt_4 = True
-    supports_stream = True
     supports_system_message = True
     supports_message_history = True
     
     default_model = 'gpt-3.5-turbo'
-    models = [
-        # Text models
+    text_models = [
         'gpt-4', 'gpt-4-0613', 'gpt-4-32k', 'gpt-4-0314', 'gpt-4-32k-0314',
         'gpt-3.5-turbo', 'gpt-3.5-turbo-16k', 'gpt-3.5-turbo-0613', 'gpt-3.5-turbo-16k-0613', 'gpt-3.5-turbo-0301',
         'gpt-3', 'text-davinci-003', 'text-davinci-002', 'code-davinci-002',
         'text-curie-001', 'text-babbage-001', 'text-ada-001',
         'davinci', 'curie', 'babbage', 'ada', 'babbage-002', 'davinci-002',
-        # Image models
-        'dalle', 'dalle-mini', 'emi'
     ]
-    
-    image_models = {"dalle", "dalle-mini", "emi"}
-    text_models = set(models) - image_models
+    image_models = ['dalle', 'dalle2', 'dalle-mini', 'emi']
+    models = [*text_models, *image_models]
     
     model_aliases = {
         "gpt-4": "gpt-4-0613",
@@ -60,16 +52,21 @@ class Nexra(AsyncGeneratorProvider, ProviderModelMixin):
         "gpt-3": "ada",
         "gpt-3": "babbage-002",
         "gpt-3": "davinci-002",
+        
+        "dalle-2": "dalle2",
     }
-
+    
+    
     @classmethod
     def get_model(cls, model: str) -> str:
-        if model in cls.models:
+        if model in cls.text_models or model in cls.image_models:
             return model
         elif model in cls.model_aliases:
             return cls.model_aliases[model]
+        elif model in cls.image_models:
+            return cls.default_image_model
         else:
-            return cls.default_model
+            return cls.default_chat_model
 
     @classmethod
     async def create_async_generator(
@@ -78,104 +75,45 @@ class Nexra(AsyncGeneratorProvider, ProviderModelMixin):
         messages: Messages,
         proxy: str = None,
         **kwargs
-    ) -> AsyncGenerator[str | ImageResponse, None]:
+    ) -> AsyncResult:
         model = cls.get_model(model)
         
-        if model in cls.image_models:
-            async for result in cls.create_image_async_generator(model, messages, proxy, **kwargs):
-                yield result
-        else:
-            async for result in cls.create_text_async_generator(model, messages, proxy, **kwargs):
-                yield result
-
-    @classmethod
-    async def create_text_async_generator(
-        cls,
-        model: str,
-        messages: Messages,
-        proxy: str = None,
-        **kwargs
-    ) -> AsyncGenerator[str, None]:
         headers = {
             "Content-Type": "application/json",
         }
+        
         async with ClientSession(headers=headers) as session:
-            data = {
-                "messages": messages,
-                "prompt": format_prompt(messages),
-                "model": model,
-                "markdown": False,
-                "stream": False,
-            }
-            async with session.post(cls.api_endpoint_text, json=data, proxy=proxy) as response:
-                response.raise_for_status()
-                result = await response.text()
-                json_result = json.loads(result)
-                yield json_result["gpt"]
-
-    @classmethod
-    async def create_image_async_generator(
-        cls,
-        model: str,
-        messages: Messages,
-        proxy: str = None,
-        **kwargs
-    ) -> AsyncGenerator[ImageResponse | str, None]:
-        headers = {
-            "Content-Type": "application/json"
-        }
-
-        prompt = messages[-1]['content'] if messages else ""
-
-        data = {
-            "prompt": prompt,
-            "model": model
-        }
-
-        async def process_response(response_text: str) -> ImageResponse | None:
-            json_start = response_text.find('{')
-            if json_start != -1:
-                json_data = response_text[json_start:]
-                try:
-                    response_data = json.loads(json_data)
-                    image_data = response_data.get('images', [])[0]
+            if model in cls.image_models:
+                # Image generation
+                prompt = messages[-1]['content'] if messages else ""
+                data = {
+                    "prompt": prompt,
+                    "model": model,
+                    "response": "url"
+                }
+                async with session.post(cls.image_api_endpoint, json=data, proxy=proxy) as response:
+                    response.raise_for_status()
+                    result = await response.text()
+                    result_json = json.loads(result.strip('_'))
+                    image_url = result_json['images'][0] if result_json['images'] else None
                     
-                    if image_data.startswith('data:image/'):
-                        return ImageResponse([image_data], "Generated image")
+                    if image_url:
+                        yield ImageResponse(images=image_url, alt=prompt)
+            else:
+                # Text completion
+                data = {
+                    "messages": messages,
+                    "prompt": format_prompt(messages),
+                    "model": model,
+                    "markdown": False
+                }
+                async with session.post(cls.chat_api_endpoint, json=data, proxy=proxy) as response:
+                    response.raise_for_status()
+                    result = await response.text()
                     
                     try:
-                        base64.b64decode(image_data)
-                        data_uri = f"data:image/jpeg;base64,{image_data}"
-                        return ImageResponse([data_uri], "Generated image")
-                    except:
-                        print("Invalid base64 data")
-                        return None
-                except json.JSONDecodeError:
-                    print("Failed to parse JSON.")
-            else:
-                print("No JSON data found in the response.")
-            return None
-
-        async with ClientSession(headers=headers) as session:
-            async with session.post(cls.api_endpoint_image, json=data, proxy=proxy) as response:
-                response.raise_for_status()
-                response_text = await response.text()
-                
-                image_response = await process_response(response_text)
-                if image_response:
-                    yield image_response
-                else:
-                    yield "Failed to process image data."
-
-    @classmethod
-    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
+                        json_response = json.loads(result)
+                        gpt_response = json_response.get('gpt', '')
+                        yield gpt_response
+                    except json.JSONDecodeError:
+                        yield result