Update api/utils.py

api/utils.py

@@ -1,7 +1,17 @@
+import uuid
+from typing import List, Dict
+from aiohttp import ClientSession
+from api.models import ChatRequest
+from api.helper import format_prompt
+from api.logger import logger
+from api.config import MODEL_MAPPING, EDITEA_API_ENDPOINT, EDITEA_HEADERS
+from fastapi import HTTPException
+
+# Editee class is now integrated within utils.py
 class Editee:
     label = "Editee"
     url = "https://editee.com"
-    api_endpoint = "https://editee.com/submit/chatgptfree"
+    api_endpoint = EDITEA_API_ENDPOINT
     working = True
     supports_stream = True
     supports_system_message = True
@@ -36,24 +46,7 @@ class Editee:
     ):
         model = cls.get_model(model)
 
-        headers = {
-            "Accept": "application/json, text/plain, */*",
-            "Accept-Language": "en-US,en;q=0.9",
-            "Cache-Control": "no-cache",
-            "Content-Type": "application/json",
-            "Origin": cls.url,
-            "Pragma": "no-cache",
-            "Priority": "u=1, i",
-            "Referer": f"{cls.url}/chat-gpt",
-            "Sec-CH-UA": '"Chromium";v="129", "Not=A?Brand";v="8"',
-            "Sec-CH-UA-Mobile": '?0',
-            "Sec-CH-UA-Platform": '"Linux"',
-            "Sec-Fetch-Dest": 'empty',
-            "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/129.0.0.0 Safari/537.36',
-            "X-Requested-With": 'XMLHttpRequest',
-        }
+        headers = EDITEA_HEADERS
 
         async with ClientSession(headers=headers) as session:
             prompt = format_prompt(messages)
@@ -63,7 +56,58 @@ class Editee:
                 "template_id": "",
                 "selected_model": model
             }
-            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
-                response.raise_for_status()
-                response_data = await response.json()
-                yield response_data['text']
+            try:
+                async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+                    response.raise_for_status()
+                    if response.content_type == 'text/event-stream':
+                        async for line in response.content:
+                            yield line.decode('utf-8')
+                    else:
+                        response_data = await response.json()
+                        yield response_data['text']
+            except Exception as e:
+                logger.error(f"Error in Editee API call: {e}")
+                raise HTTPException(status_code=500, detail="Error in Editee API call")
+
+# Function to process the response
+async def process_response(request: ChatRequest, stream: bool = False):
+    try:
+        model = MODEL_MAPPING.get(request.model, request.model)
+        messages = [
+            {"role": message.role, "content": message.content}
+            for message in request.messages
+        ]
+
+        generator = Editee.create_async_generator(
+            model=model,
+            messages=messages,
+            proxy=None  # Add proxy if needed
+        )
+
+        if stream:
+            async def event_generator():
+                async for chunk in generator:
+                    yield f"data: {chunk}\n\n"
+            return event_generator()
+        else:
+            full_response = ""
+            async for chunk in generator:
+                full_response += chunk
+
+            return {
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion",
+                "created": int(uuid.uuid1().time),
+                "model": model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": full_response},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": None,
+            }
+    except Exception as e:
+        logger.error(f"Error processing response: {e}")
+        raise HTTPException(status_code=500, detail=str(e))