Update main.py

main.py

@@ -10,6 +10,7 @@ from pathlib import Path
 import json
 import datetime
 import time
+import threading
 from typing import Optional, Dict, List, Any, Generator
 import asyncio
 from starlette.status import HTTP_403_FORBIDDEN
@@ -380,57 +381,106 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
     else:
         endpoint = env_vars['secret_api_endpoint']
         custom_headers = {}
-    print(payload_dict)
-    print(f"laude ye endpoint use kia: {endpoint}")
-    scraper = get_scraper()
     
-    async def stream_generator(payload_dict):
+    print(f"Using endpoint: {endpoint}")
+    
+    # Create a new scraper for each request to avoid potential blocking
+    scraper = cloudscraper.create_scraper(browser={
+        'browser': 'chrome',
+        'platform': 'windows',
+        'mobile': False
+    })
+    
+    # Set a timeout for the entire request handling
+    TIMEOUT_SECONDS = 20 
+    
+    async def stream_generator_with_timeout(payload_dict):
         try:
-            # Send POST request with the correct headers
-            response = scraper.post(
-                f"{endpoint}/v1/chat/completions",
-                json=payload_dict,
-                headers=custom_headers,
-                stream=True
-            )
-
-            # Handle response errors
-            if response.status_code >= 400:
-                error_messages = {
-                    422: "Unprocessable entity. Check your payload.",
-                    400: "Bad request. Verify input data.",
-                    403: "Forbidden. You do not have access to this resource.",
-                    404: "The requested resource was not found.",
-                }
-                detail = error_messages.get(response.status_code, f"Error code: {response.status_code}")
-                raise HTTPException(status_code=response.status_code, detail=detail)
-
-            # Stream response lines to the client - use buffer for efficiency
-            buffer = []
-            buffer_size = 0
-            max_buffer = 8192  # 8KB buffer
+            # Create a thread-safe event for cancellation
+            cancel_event = threading.Event()
             
-            for line in response.iter_lines():
-                if line:
-                    decoded = line.decode('utf-8') + "\n"
-                    buffer.append(decoded)
-                    buffer_size += len(decoded)
+            def request_with_timeout():
+                try:
+                    # Send POST request with the correct headers and timeout
+                    response = scraper.post(
+                        f"{endpoint}/v1/chat/completions",
+                        json=payload_dict,
+                        headers=custom_headers,
+                        stream=True,
+                        timeout=TIMEOUT_SECONDS
+                    )
                     
-                    if buffer_size >= max_buffer:
-                        yield ''.join(buffer)
-                        buffer = []
-                        buffer_size = 0
+                    # Handle response errors
+                    if response.status_code >= 400:
+                        error_messages = {
+                            422: "Unprocessable entity. Check your payload.",
+                            400: "Bad request. Verify input data.",
+                            403: "Forbidden. You do not have access to this resource.",
+                            404: "The requested resource was not found.",
+                        }
+                        detail = error_messages.get(response.status_code, f"Error code: {response.status_code}")
+                        return {"error": detail, "status_code": response.status_code}
+                    
+                    result = []
+                    
+                    # Process the streaming response with timeout checks
+                    for line in response.iter_lines():
+                        # Check for cancellation
+                        if cancel_event.is_set():
+                            break
+                            
+                        if line:
+                            decoded = line.decode('utf-8') + "\n"
+                            result.append(decoded)
+                    
+                    return {"lines": result}
+                    
+                except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
+                    return {"error": "Request timed out or connection failed", "status_code": 504}
+                except Exception as e:
+                    return {"error": str(e), "status_code": 500}
             
-            # Flush remaining buffer
-            if buffer:
-                yield ''.join(buffer)
-
+            # Execute request in a ThreadPoolExecutor with a timeout
+            loop = asyncio.get_running_loop()
+            with ThreadPoolExecutor() as pool:
+                response_future = loop.run_in_executor(pool, request_with_timeout)
+                
+                try:
+                    # Wait for response with a timeout
+                    response_data = await asyncio.wait_for(response_future, timeout=TIMEOUT_SECONDS)
+                    
+                    # If there was an error, raise an HTTPException
+                    if "error" in response_data:
+                        raise HTTPException(
+                            status_code=response_data.get("status_code", 500),
+                            detail=response_data["error"]
+                        )
+                    
+                    # Stream the response lines
+                    for line in response_data.get("lines", []):
+                        yield line
+                        
+                except asyncio.TimeoutError:
+                    # Cancel the ongoing request
+                    cancel_event.set()
+                    raise HTTPException(status_code=504, detail="Request timed out after 20 seconds")
+                    
         except Exception as e:
+            if isinstance(e, HTTPException):
+                raise e
             # Use a generic error message that doesn't expose internal details
-            raise HTTPException(status_code=500, detail="An error occurred while processing your request")
-
-    return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
+            raise HTTPException(status_code=500, detail=f"An error occurred while processing your request: {str(e)}")
 
+    # Return streaming response with proper timeout handling
+    try:
+        return StreamingResponse(
+            stream_generator_with_timeout(payload_dict), 
+            media_type="application/json"
+        )
+    except Exception as e:
+        if isinstance(e, HTTPException):
+            raise e
+        raise HTTPException(status_code=500, detail=f"Failed to initialize streaming response: {str(e)}")
 # Asynchronous logging function
 async def log_request(request, model):
     # Get minimal data for logging