Update api/utils.py

api/utils.py

@@ -1,11 +1,11 @@
+# api/utils.py
+
 from datetime import datetime
 import json
 from typing import AsyncGenerator, Union
-import uuid
 
 import aiohttp
 from fastapi import HTTPException
-from fastapi.responses import JSONResponse
 from api.config import GIZAI_API_ENDPOINT, GIZAI_BASE_URL
 from api.models import ChatRequest, ImageResponseModel, ChatCompletionResponse
 from api.logger import setup_logger
@@ -68,10 +68,10 @@ class GizAI:
     def is_image_model(cls, model: str) -> bool:
         return model in cls.image_models
 
-async def process_gizai_response(request: ChatRequest, model: str) -> Union[AsyncGenerator[str, None], JSONResponse]:
+async def process_gizai_response(request: ChatRequest, model: str) -> Union[AsyncGenerator[str, None], ImageResponseModel, ChatCompletionResponse]:
     async with aiohttp.ClientSession() as session:
         if GizAI.is_image_model(model):
-            # Image generation
+            # Image generation logic
             prompt = request.messages[-1].content if isinstance(request.messages[-1].content, str) else request.messages[-1].content[0].get("text", "")
             data = {
                 "model": model,
@@ -110,7 +110,7 @@ async def process_gizai_response(request: ChatRequest, model: str) -> Union[Asyn
                     response_data = await response.json()
                     if response_data.get('status') == 'completed' and response_data.get('output'):
                         images = response_data['output']
-                        return {"images": images, "alt": "Generated Image"}
+                        return ImageResponseModel(images=images, alt="Generated Image")
                     else:
                         raise HTTPException(status_code=500, detail="Image generation failed.")
             except aiohttp.ClientResponseError as e:
@@ -120,7 +120,7 @@ async def process_gizai_response(request: ChatRequest, model: str) -> Union[Asyn
                 logger.error(f"Unexpected error: {str(e)}")
                 raise HTTPException(status_code=500, detail=str(e))
         else:
-            # Chat completion
+            # Chat completion logic
             messages_formatted = [
                 {
                     "type": "human",
@@ -170,7 +170,20 @@ async def process_gizai_response(request: ChatRequest, model: str) -> Union[Asyn
                     else:
                         # Handle non-streaming response
                         result = await response.json()
-                        return result.get('output', '')
+                        return ChatCompletionResponse(
+                            id=f"chatcmpl-{uuid.uuid4()}",
+                            object="chat.completion",
+                            created=int(datetime.now().timestamp()),
+                            model=model,
+                            choices=[
+                                {
+                                    "index": 0,
+                                    "message": {"role": "assistant", "content": result.get('output', '')},
+                                    "finish_reason": "stop",
+                                }
+                            ],
+                            usage=None,
+                        )
             except aiohttp.ClientResponseError as e:
                 logger.error(f"HTTP error occurred: {e.status} - {e.message}")
                 raise HTTPException(status_code=e.status, detail=str(e))