Update main.py

main.py

@@ -6,9 +6,8 @@ from datetime import datetime
 from typing import Any, Dict, List, Optional
 
 import httpx
-import uvicorn
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Depends, Request
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
@@ -176,13 +175,18 @@ async def chat_completions(
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         if line and ("[DONE]" not in line):
-                            # Assuming the response starts with some prefix, adjust accordingly
                             try:
-                                data = json.loads(line[5:]).get("data", {})
-                                content = data.get("text", "")
-                                yield f"data: {json.dumps(create_chat_completion_data(content, request.model))}\n\n"
+                                # Adjust the slicing based on the actual response format
+                                data = json.loads(line[5:]) if len(line) > 5 else None
+                                if data and "data" in data and "text" in data["data"]:
+                                    content = data["data"].get("text", "")
+                                    yield f"data: {json.dumps(create_chat_completion_data(content, request.model))}\n\n"
+                                else:
+                                    logger.warning(f"Unexpected data format: {line}")
                             except json.JSONDecodeError as e:
-                                logger.error(f"JSON decode error: {e}")
+                                logger.error(f"JSON decode error for line: {line} - {e}")
+                            except Exception as e:
+                                logger.error(f"Error processing line: {line} - {e}")
                     # Indicate the end of the stream
                     yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
                     yield "data: [DONE]\n\n"
@@ -190,7 +194,7 @@ async def chat_completions(
                 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"Request error occurred: {e}")
+                logger.error(f"An error occurred while requesting: {e}")
                 raise HTTPException(status_code=500, detail=str(e))
 
     if request.stream:
@@ -201,9 +205,14 @@ async def chat_completions(
         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"]
+                try:
+                    data = json.loads(chunk[6:])
+                    if data["choices"][0]["delta"].get("content"):
+                        full_response += data["choices"][0]["delta"]["content"]
+                except json.JSONDecodeError as e:
+                    logger.error(f"JSON decode error in non-streaming response: {chunk} - {e}")
+                except Exception as e:
+                    logger.error(f"Error processing chunk in non-streaming response: {chunk} - {e}")
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
@@ -219,7 +228,3 @@ async def chat_completions(
             ],
             "usage": None,
         }
-
-# Remove the Uvicorn run block for production deployment
-# if __name__ == "__main__":
-#     uvicorn.run(app, host="0.0.0.0", port=7860)