Update main.py

main.py

@@ -14,16 +14,22 @@ from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
+# Configure logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
+# Load environment variables from .env file
 load_dotenv()
+
+# Initialize FastAPI app
 app = FastAPI()
+
+# Constants and configurations
 BASE_URL = "https://aichatonlineorg.erweima.ai/aichatonline"
-APP_SECRET = os.getenv("APP_SECRET","666")
-ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN","")
+APP_SECRET = os.getenv("APP_SECRET", "666")
+ACCESS_TOKEN = os.getenv("SD_ACCESS_TOKEN", "")
 headers = {    
     'accept': '*/*',
     'accept-language': 'en-US,en;q=0.9',
@@ -38,6 +44,7 @@ headers = {
     'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36', 
 }
 
+# Define allowed models
 ALLOWED_MODELS = [
     {"id": "claude-3.5-sonnet", "name": "claude-3.5-sonnet"},
     {"id": "claude-3-opus", "name": "claude-3-opus"},
@@ -47,28 +54,30 @@ ALLOWED_MODELS = [
     {"id": "o1-mini", "name": "o1-mini"},
     {"id": "gpt-4o-mini", "name": "gpt-4o-mini"},
 ]
-# Configure CORS
+
+# Configure CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow all sources, you can restrict specific sources if needed
+    allow_origins=["*"],  # Allow all origins; restrict if necessary
     allow_credentials=True,
-    allow_methods=["*"],  # All methods allowed
+    allow_methods=["*"],  # Allow all HTTP methods
     allow_headers=["*"],  # Allow all headers
 )
-security = HTTPBearer()
 
+# Security configuration
+security = HTTPBearer()
 
+# Pydantic models
 class Message(BaseModel):
     role: str
     content: str
 
-
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
 
-
+# Helper functions
 def simulate_data(content, model):
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -85,7 +94,6 @@ def simulate_data(content, model):
         "usage": None,
     }
 
-
 def stop_data(content, model):
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -101,8 +109,7 @@ def stop_data(content, model):
         ],
         "usage": None,
     }
-    
-    
+
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -119,12 +126,16 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
         "usage": None,
     }
 
-
 def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
     if credentials.credentials != APP_SECRET:
         raise HTTPException(status_code=403, detail="Invalid APP_SECRET")
     return credentials.credentials
 
+# Utility function to replace escaped newlines
+def replace_escaped_newlines(input_string: str) -> str:
+    return input_string.replace("\\n", "\n")
+
+# API Endpoints
 
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
@@ -137,72 +148,68 @@ async def chat_completions_options():
         },
     )
 
-
-def replace_escaped_newlines(input_string: str) -> str:
-    return input_string.replace("\\n", "\n")
-
-
 @app.get("/hf/v1/models")
 async def list_models():
     return {"object": "list", "data": ALLOWED_MODELS}
 
-
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(
     request: ChatRequest, app_secret: str = Depends(verify_app_secret)
 ):
     logger.info(f"Received chat completion request for model: {request.model}")
 
+    # Validate 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)}",
         )
+
     # Generate a UUID
     original_uuid = uuid.uuid4()
     uuid_str = str(original_uuid).replace("-", "")
 
-json_data = {
-    'prompt': "\n".join(
-        [
-            f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
-            for msg in request.messages
-        ]
-    ),
-    'stream': True,
-    'app_name': 'ChitChat_Edge_Ext',
-    'app_version': '4.28.0',
-    'tz_name': 'Asia/Karachi',
-    'cid': 'C092SEMXM9BJ',
-    'model': request.model,
-    'search': False,  # Ensure search is disabled
-    'auto_search': False,  # Ensure auto_search is disabled
-    'filter_search_history': False,
-    'from': 'chat',
-    'group_id': 'default',
-    'chat_models': [],
-    'files': [],
-    'prompt_template': {
-        'key': '',
-        'attributes': {
-            'lang': 'original',
+    # Construct the payload to send to the external API
+    json_data = {
+        'prompt': "\n".join(
+            [
+                f"{'User' if msg.role == 'user' else 'Assistant'}: {msg.content}"
+                for msg in request.messages
+            ]
+        ),
+        'stream': True,
+        'app_name': 'ChitChat_Edge_Ext',
+        'app_version': '4.28.0',
+        'tz_name': 'Asia/Karachi',
+        'cid': 'C092SEMXM9BJ',
+        'model': request.model,
+        'search': False,  # Ensure search is disabled
+        'auto_search': False,  # Ensure auto_search is disabled
+        'filter_search_history': False,
+        'from': 'chat',
+        'group_id': 'default',
+        'chat_models': [],
+        'files': [],
+        'prompt_template': {
+            'key': '',
+            'attributes': {
+                'lang': 'original',
+            },
         },
-    },
-    'tools': {
-        'auto': [
-            'text_to_image',
-            'data_analysis',
-            # Removed 'search' from the list
-        ],
-    },
-    'extra_info': {
-        'origin_url': '',
-        'origin_title': '',
-    },
-}
-
-
+        'tools': {
+            'auto': [
+                'text_to_image',
+                'data_analysis',
+                # 'search' has been removed to disable search functionality
+            ],
+        },
+        'extra_info': {
+            'origin_url': '',
+            'origin_title': '',
+        },
+    }
 
+    # Define the asynchronous generator for streaming responses
     async def generate():
         async with httpx.AsyncClient() as client:
             try:
@@ -210,10 +217,28 @@ json_data = {
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         if line and ("[DONE]" not in line):
-                            content = json.loads(line[5:])["data"]
-                            yield f"data: {json.dumps(create_chat_completion_data(content.get('text',''), request.model))}\n\n"
-                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
-                    yield "data: [DONE]\n\n"
+                            # Assuming the line starts with some prefix before JSON, e.g., "data: "
+                            # Adjust if necessary based on actual response format
+                            try:
+                                # Remove any prefix before JSON if present
+                                if line.startswith("data: "):
+                                    line_content = line[6:]
+                                else:
+                                    line_content = line
+                                
+                                # Parse the JSON content
+                                content = json.loads(line_content)["data"]
+                                
+                                # Yield the formatted data
+                                yield f"data: {json.dumps(create_chat_completion_data(content.get('text',''), request.model))}\n\n"
+                            except json.JSONDecodeError as e:
+                                logger.error(f"JSON decode error: {e}")
+                                continue
+                        else:
+                            # Signal the end of the stream
+                            if line and "[DONE]" in line:
+                                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))
@@ -229,11 +254,14 @@ json_data = {
         full_response = ""
         async for chunk in generate():
             if chunk.startswith("data: ") and not chunk[6:].startswith("[DONE]"):
-                # print(chunk)
-                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: {e}")
+                    continue
+
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
@@ -249,7 +277,6 @@ json_data = {
             "usage": None,
         }
 
-
-
+# Entry point for running the application
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)