Update api/utils.py

api/utils.py

@@ -59,18 +59,20 @@ class Editee:
             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:
+                    # Check if the response is in JSON format
+                    if response.headers.get('Content-Type') == 'application/json':
                         response_data = await response.json()
                         yield response_data['text']
+                    else:
+                        # Stream the response line by line
+                        async for line in response.content:
+                            yield line.decode('utf-8')
             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):
+# Function to process non-streaming response
+async def process_response(request: ChatRequest):
     try:
         model = MODEL_MAPPING.get(request.model, request.model)
         messages = [
@@ -84,30 +86,52 @@ async def process_response(request: ChatRequest, stream: bool = False):
             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
+        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,
-            }
+        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))
+
+# Function to process streaming response
+def process_response_stream(request: ChatRequest):
+    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
+        )
+
+        async def event_generator():
+            try:
+                async for chunk in generator:
+                    yield f"data: {chunk}\n\n"
+                yield "data: [DONE]\n\n"
+            except Exception as e:
+                logger.error(f"Error in streaming response: {e}")
+                raise HTTPException(status_code=500, detail=str(e))
+        return event_generator()
+    except Exception as e:
+        logger.error(f"Error processing streaming response: {e}")
+        raise HTTPException(status_code=500, detail=str(e))