Update main.py

main.py

@@ -14,9 +14,14 @@ from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
+# ==============================
+# Configuration and Setup
+# ==============================
+
 # Configure logging
 logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    level=logging.DEBUG,  # Set to DEBUG for detailed logs
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
@@ -67,7 +72,10 @@ app.add_middleware(
 # Security configuration
 security = HTTPBearer()
 
-# Pydantic models
+# ==============================
+# Pydantic Models
+# ==============================
+
 class Message(BaseModel):
     role: str
     content: str
@@ -77,8 +85,12 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     stream: Optional[bool] = False
 
-# Helper functions
-def simulate_data(content, model):
+# ==============================
+# Helper Functions
+# ==============================
+
+def simulate_data(content: str, model: str) -> Dict[str, Any]:
+    """Simulate chunked response data."""
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -94,7 +106,8 @@ def simulate_data(content, model):
         "usage": None,
     }
 
-def stop_data(content, model):
+def stop_data(content: str, model: str) -> Dict[str, Any]:
+    """Simulate the final chunk indicating the end of the response."""
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -111,6 +124,7 @@ def stop_data(content, model):
     }
 
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
+    """Create a structured chat completion data chunk."""
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -126,19 +140,23 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
         "usage": None,
     }
 
-def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
+def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
+    """Verify the provided APP_SECRET."""
     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:
+    """Replace escaped newline characters with actual newlines."""
     return input_string.replace("\\n", "\n")
 
+# ==============================
 # API Endpoints
+# ==============================
 
 @app.options("/hf/v1/chat/completions")
 async def chat_completions_options():
+    """Handle CORS preflight requests."""
     return Response(
         status_code=200,
         headers={
@@ -150,19 +168,22 @@ async def chat_completions_options():
 
 @app.get("/hf/v1/models")
 async def list_models():
+    """List all allowed 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)
 ):
+    """Handle chat completion requests."""
     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]:
+    allowed_model_ids = [model['id'] for model in ALLOWED_MODELS]
+    if request.model not in allowed_model_ids:
         raise HTTPException(
             status_code=400,
-            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(allowed_model_ids)}",
         )
 
     # Generate a UUID
@@ -183,8 +204,8 @@ async def chat_completions(
         'tz_name': 'Asia/Karachi',
         'cid': 'C092SEMXM9BJ',
         'model': request.model,
-        'search': False,  # Ensure search is disabled
-        'auto_search': False,  # Ensure auto_search is disabled
+        'search': False,  # Disable search
+        'auto_search': False,  # Disable auto_search
         'filter_search_history': False,
         'from': 'chat',
         'group_id': 'default',
@@ -198,9 +219,9 @@ async def chat_completions(
         },
         'tools': {
             'auto': [
+                'search',  # Re-add search to maintain API expectations
                 'text_to_image',
                 'data_analysis',
-                # 'search' has been removed to disable search functionality
             ],
         },
         'extra_info': {
@@ -216,26 +237,43 @@ async def chat_completions(
                 async with client.stream('POST', 'https://sider.ai/api/v3/completion/text', headers=headers, json=json_data, timeout=120.0) as response:
                     response.raise_for_status()
                     async for line in response.aiter_lines():
+                        if line:
+                            logger.debug(f"Raw line received: {line}")  # Log raw line
                         if line and ("[DONE]" not in line):
-                            # 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
+                                # Remove 'data: ' prefix if present
                                 if line.startswith("data: "):
                                     line_content = line[6:]
                                 else:
                                     line_content = line
                                 
-                                # Parse the JSON content
-                                content = json.loads(line_content)["data"]
+                                # Log the content before parsing
+                                logger.debug(f"Line content to parse: {line_content}")
+                                
+                                # Check if the line is not empty after stripping
+                                if not line_content.strip():
+                                    logger.warning("Received an empty line, skipping.")
+                                    continue
+                                
+                                # Attempt to parse JSON
+                                parsed_json = json.loads(line_content)
+                                
+                                # Ensure 'data' key exists
+                                if "data" not in parsed_json:
+                                    logger.error(f"'data' key not found in the response: {parsed_json}")
+                                    continue
+                                
+                                content_data = parsed_json["data"]
+                                
+                                # Extract text content if available
+                                text_content = content_data.get("text", "")
                                 
                                 # Yield the formatted 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(text_content, request.model))}\n\n"
                             except json.JSONDecodeError as e:
-                                logger.error(f"JSON decode error: {e}")
+                                logger.error(f"JSON decode error: {e} | Line: {line_content}")
                                 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"
@@ -277,6 +315,9 @@ async def chat_completions(
             "usage": None,
         }
 
-# Entry point for running the application
+# ==============================
+# Entry Point
+# ==============================
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)