Update main.py

main.py

@@ -187,28 +187,74 @@ class ChatRequest(BaseModel):
 from fastapi.responses import Response
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatRequest):
-    messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
-
-    async_generator = Blackbox.create_async_generator(
-        model=request.model,
-        messages=messages
-    )
-
-    response_content = ""
-    async for chunk in async_generator:
-        if isinstance(chunk, str):
-            response_content += chunk
-        else:
-            response_content += chunk.content
+async def chat_completions(
+    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
+):
+    logger.info(f"Received chat completion request for model: {request.model}")
+
+    if request.model not in [model['id'] for model in ALLOWED_MODELS]:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+        )
 
-    # Clean up the response to extract the plain text content
-    clean_content = response_content.replace('\\n', '\n').strip()
+    # Generate a UUID for the conversation
+    conversation_id = str(uuid.uuid4()).replace("-", "")
 
-    return {
-        "id": str(uuid.uuid4()),
-        "object": "chat.completion",
-        "created": int(datetime.now().timestamp()),
-        "model": request.model,
-        "choices": [{"index": 0, "message": {"role": "assistant", "content": clean_content}, "finish_reason": "stop"}]
+    json_data = {
+        "attachments": [],
+        "conversationId": conversation_id,
+        "prompt": "\n".join(
+            [
+                f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
+                for msg in request.messages
+            ]
+        ),
     }
+
+    headers["uniqueid"] = conversation_id
+
+    async def generate():
+        async with httpx.AsyncClient() as client:
+            try:
+                async with client.stream('POST', f'{BASE_URL}/api/chat/gpt4o/chat', headers=headers, json=json_data, timeout=120.0) as response:
+                    response.raise_for_status()
+                    async for line in response.aiter_lines():
+                        if line and line != "[DONE]":
+                            content = json.loads(line)["data"]
+                            yield f"data: {json.dumps(create_chat_completion_data(content['message'], request.model))}\n\n"
+                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
+                    yield "data: [DONE]\n\n"
+            except httpx.HTTPStatusError as e:
+                logger.error(f"HTTP error occurred: {e}")
+                raise HTTPException(status_code=e.response.status_code, detail=str(e))
+            except httpx.RequestError as e:
+                logger.error(f"An error occurred while requesting: {e}")
+                raise HTTPException(status_code=500, detail=str(e))
+
+    if request.stream:
+        logger.info("Streaming response")
+        return StreamingResponse(generate(), media_type="text/event-stream")
+    else:
+        logger.info("Non-streaming response")
+        full_response = ""
+        async for chunk in generate():
+            if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
+                data = json.loads(chunk[6:])
+                if data["choices"][0]["delta"].get("content"):
+                    full_response += data["choices"][0]["delta"]["content"]
+
+        return {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": full_response},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": None,
+        }