Update main.py

main.py

@@ -455,6 +455,18 @@ class Blackbox:
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
+        headers_chat = {
+            'Accept': 'text/x-component',
+            'Content-Type': 'text/plain;charset=UTF-8',
+            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
+            'next-action': next_action,
+            'next-router-state-tree': next_router_state_tree,
+            'next-url': '/'
+        }
+        headers_chat_combined = {**common_headers, **headers_chat}
+
+        data_chat = '[]'
+
         async with ClientSession(headers=common_headers, timeout=ClientTimeout(total=60)) as session:
             try:
                 async with session.post(
@@ -610,17 +622,25 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
         if request.stream:
             # Streaming response
-            async def content_generator():
+            async def event_generator():
                 async for chunk in Blackbox.create_async_generator(
                     model=request.model,
                     messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
                     proxy=None,  # Add proxy if needed
                     websearch=False  # Modify if websearch is needed
                 ):
-                    yield chunk
+                    if isinstance(chunk, ImageResponseModel):
+                        # Yield image URLs as plain text in SSE format
+                        yield f"data: Image URL: {chunk.images}\n\n"
+                    else:
+                        # Ensure chunk is a string and yield as plain text in SSE format
+                        if isinstance(chunk, str):
+                            yield f"data: {chunk}\n\n"
+                        else:
+                            yield f"data: {str(chunk)}\n\n"
 
             logger.info(f"Initiating streaming response for API key: {api_key} | IP: {client_ip}")
-            return StreamingResponse(content_generator(), media_type='text/plain')
+            return StreamingResponse(event_generator(), media_type='text/event-stream')
         else:
             # Non-streaming response
             response_content = await Blackbox.generate_response(
@@ -632,7 +652,25 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
             logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
             return {
-                "content": response_content
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion",
+                "created": int(datetime.now().timestamp()),
+                "model": request.model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": response_content
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
+                    "completion_tokens": len(response_content.split()),
+                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
+                },
             }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")