Update api/utils.py

api/utils.py

@@ -68,8 +68,87 @@ 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], ImageResponseModel, ChatCompletionResponse]:
+async def process_gizai_stream_response(request: ChatRequest, model: str) -> AsyncGenerator[str, None]:
     async with aiohttp.ClientSession() as session:
+        # Set up headers
+        headers = {
+            'Accept': 'application/json, text/plain, */*',
+            'Accept-Language': 'en-US,en;q=0.9',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'Content-Type': 'application/json',
+            'Origin': 'https://app.giz.ai',
+            'Pragma': 'no-cache',
+            '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/130.0.0.0 Safari/537.36',
+            'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"'
+        }
+
+        if GizAI.is_image_model(model):
+            # Image generation logic (streaming might not make sense here)
+            raise HTTPException(status_code=400, detail="Image generation does not support streaming.")
+        else:
+            # Chat completion logic
+            messages_formatted = [
+                {
+                    "type": "human",
+                    "content": msg.content if isinstance(msg.content, str) else msg.content[0].get("text", "")
+                } for msg in request.messages
+            ]
+            data = {
+                "model": model,
+                "input": {
+                    "messages": messages_formatted,
+                    "mode": "plan"
+                },
+                "noStream": False  # Enable streaming
+            }
+            try:
+                async with session.post(
+                    GIZAI_API_ENDPOINT,
+                    headers=headers,
+                    json=data
+                ) as response:
+                    response.raise_for_status()
+                    async for line in response.content:
+                        if line:
+                            decoded_line = line.decode('utf-8').strip()
+                            if decoded_line.startswith("data:"):
+                                content = decoded_line.replace("data: ", "")
+                                yield f"data: {content}\n\n"
+                # Indicate the end of the stream
+                yield "data: [DONE]\n\n"
+            except aiohttp.ClientResponseError as e:
+                logger.error(f"HTTP error occurred: {e.status} - {e.message}")
+                raise HTTPException(status_code=e.status, detail=str(e))
+            except Exception as e:
+                logger.error(f"Unexpected error: {str(e)}")
+                raise HTTPException(status_code=500, detail=str(e))
+
+async def process_gizai_non_stream_response(request: ChatRequest, model: str) -> Union[ImageResponseModel, ChatCompletionResponse]:
+    async with aiohttp.ClientSession() as session:
+        # Set up headers
+        headers = {
+            'Accept': 'application/json, text/plain, */*',
+            'Accept-Language': 'en-US,en;q=0.9',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'Content-Type': 'application/json',
+            'Origin': 'https://app.giz.ai',
+            'Pragma': 'no-cache',
+            '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/130.0.0.0 Safari/537.36',
+            'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"'
+        }
+
         if GizAI.is_image_model(model):
             # Image generation logic
             prompt = request.messages[-1].content if isinstance(request.messages[-1].content, str) else request.messages[-1].content[0].get("text", "")
@@ -88,22 +167,7 @@ async def process_gizai_response(request: ChatRequest, model: str) -> Union[Asyn
             try:
                 async with session.post(
                     GIZAI_API_ENDPOINT,
-                    headers={
-                        'Accept': 'application/json, text/plain, */*',
-                        'Accept-Language': 'en-US,en;q=0.9',
-                        'Cache-Control': 'no-cache',
-                        'Connection': 'keep-alive',
-                        'Content-Type': 'application/json',
-                        'Origin': 'https://app.giz.ai',
-                        'Pragma': 'no-cache',
-                        '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/130.0.0.0 Safari/537.36',
-                        'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
-                        'sec-ch-ua-mobile': '?0',
-                        'sec-ch-ua-platform': '"Linux"'
-                    },
+                    headers=headers,
                     json=data
                 ) as response:
                     response.raise_for_status()
@@ -133,57 +197,30 @@ async def process_gizai_response(request: ChatRequest, model: str) -> Union[Asyn
                     "messages": messages_formatted,
                     "mode": "plan"
                 },
-                "noStream": not request.stream
+                "noStream": True  # Disable streaming
             }
             try:
                 async with session.post(
                     GIZAI_API_ENDPOINT,
-                    headers={
-                        'Accept': 'application/json, text/plain, */*',
-                        'Accept-Language': 'en-US,en;q=0.9',
-                        'Cache-Control': 'no-cache',
-                        'Connection': 'keep-alive',
-                        'Content-Type': 'application/json',
-                        'Origin': 'https://app.giz.ai',
-                        'Pragma': 'no-cache',
-                        '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/130.0.0.0 Safari/537.36',
-                        'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
-                        'sec-ch-ua-mobile': '?0',
-                        'sec-ch-ua-platform': '"Linux"'
-                    },
+                    headers=headers,
                     json=data
                 ) as response:
                     response.raise_for_status()
-                    if request.stream:
-                        # Handle streaming response
-                        async def stream_response():
-                            async for line in response.content:
-                                if line:
-                                    decoded_line = line.decode('utf-8').strip()
-                                    if decoded_line.startswith("data:"):
-                                        content = decoded_line.replace("data: ", "")
-                                        yield f"data: {content}\n\n"
-                        return stream_response()
-                    else:
-                        # Handle non-streaming response
-                        result = await response.json()
-                        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,
-                        )
+                    result = await response.json()
+                    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))