Update app.py

app.py

@@ -12,20 +12,19 @@ from pydantic import BaseModel
 # Configuration
 API_CHALLENGE_URL = 'https://api.eqing.tech/api/altcaptcha/challenge'
 NEXTWAY_CHAT_URL = 'https://origin.eqing.tech/api/openai/v1/chat/completions'
-CREDENTIAL_EXPIRY_MARGIN = 60 * 1000  # 60 seconds in milliseconds
+CREDENTIAL_EXPIRY_MARGIN = 60 * 1000
 PORT = 7860
 API_ENDPOINT = '/v1/chat/completions'
 MODEL_NAME = "gpt-4o-free"
-REQUEST_TIMEOUT = 480  # seconds
+REQUEST_TIMEOUT = 480
 MAX_RETRIES = 3
-RETRY_DELAY = 1  # seconds
+RETRY_DELAY = 1
 
 # Global variables
 current_credential = None
 credential_expiry = None
 is_refreshing_credential = False
 
-# Model mappings
 MODEL_MAPPING = {
     'gpt-4o-all-lite': 'gpt-4o-mini',
     'gpt-4o-image': 'gpt-4o-mini-image-free',
@@ -45,32 +44,47 @@ class ChatRequest(BaseModel):
     top_p: Optional[float] = 1
     max_tokens: Optional[int] = 4000
 
-def extract_content(text: str) -> str:
-    lines = text.split('\n')
-    ai_response = ''
-    ignored_id = 'chatcmpl-nxaTnETiUyAVBjdRwlr38Yt3'
-    created = 1687070102
-
-    for line in lines:
-        line = line.strip()
-        if line.startswith('data:'):
-            data_str = line[5:].strip()
-            if not data_str or data_str in ['[ORIGIN]', '[DONE]']:
-                continue
-
-            try:
-                json_data = json.loads(data_str)
-                if json_data.get('id') == ignored_id or json_data.get('created') == created:
-                    continue
-
-                if (json_data.get('choices') and json_data['choices'][0].get('delta') and 
-                    'content' in json_data['choices'][0]['delta']):
-                    ai_response += json_data['choices'][0]['delta']['content']
-
-            except json.JSONDecodeError:
-                print(f'Skipping non-JSON data')
+async def extract_streaming_content(response):
+    buffer = ""
+    complete_response = ""
+    async for chunk in response.content:
+        buffer += chunk.decode()
+        lines = buffer.split('\n')
+        
+        # Process complete lines
+        for line in lines[:-1]:
+            line = line.strip()
+            if line.startswith('data: '):
+                content = process_data_line(line[5:].strip())
+                if content:
+                    complete_response += content
+        
+        # Keep the last incomplete line in buffer
+        buffer = lines[-1] if lines else ""
+    
+    # Process any remaining data
+    if buffer:
+        line = buffer.strip()
+        if line.startswith('data: '):
+            content = process_data_line(line[5:].strip())
+            if content:
+                complete_response += content
+    
+    return complete_response
 
-    return ai_response
+def process_data_line(data_str: str) -> str:
+    if not data_str or data_str in ['[ORIGIN]', '[DONE]']:
+        return ''
+    
+    try:
+        json_data = json.loads(data_str)
+        if (json_data.get('choices') and 
+            json_data['choices'][0].get('delta') and 
+            'content' in json_data['choices'][0]['delta']):
+            return json_data['choices'][0]['delta']['content']
+    except json.JSONDecodeError:
+        print(f'Invalid JSON data: {data_str}')
+    return ''
 
 async def solve_challenge(challenge: str, salt: str, algorithm: str = "SHA-512", max_number: int = 1000000):
     start_time = time.time()
@@ -155,36 +169,7 @@ async def get_credential():
     
     return current_credential
 
-@app.post(API_ENDPOINT)
-async def chat_endpoint(request: ChatRequest):
-    try:
-        model = MODEL_MAPPING.get(request.model, request.model or MODEL_NAME)
-        response_content = await handle_chat_request(model, request.messages)
-        
-        if response_content is None:
-            return Response(
-                content="Failed to get response from API",
-                status_code=500
-            )
-
-        return {
-            "choices": [{
-                "message": {
-                    "role": "assistant",
-                    "content": response_content.strip()
-                },
-                "finish_reason": "stop",
-                "index": 0
-            }],
-            "model": model,
-            "object": "chat.completion"
-        }
-
-    except Exception as e:
-        print(f"Error handling chat request: {e}")
-        return Response(content="Internal server error", status_code=500)
-
-async def handle_chat_request(model: str, messages: List[Dict[str, str]]):
+async def handle_chat_request(model: str, messages: List[Dict[str, str]], retries=0):
     captcha_token = await get_credential()
     if not captcha_token:
         return None
@@ -206,24 +191,51 @@ async def handle_chat_request(model: str, messages: List[Dict[str, str]]):
             async with session.post(
                 NEXTWAY_CHAT_URL,
                 json=body,
-                timeout=REQUEST_TIMEOUT
+                timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)
             ) as response:
                 if response.status != 200:
+                    if retries < MAX_RETRIES:
+                        await asyncio.sleep(RETRY_DELAY)
+                        return await handle_chat_request(model, messages, retries + 1)
                     return None
 
-                complete_response = ""
-                async for chunk in response.content:
-                    chunk_text = chunk.decode()
-                    content = extract_content(chunk_text)
-                    if content:
-                        complete_response += content
-
-                return complete_response
+                return await extract_streaming_content(response)
 
-        except Exception as e:
+        except (aiohttp.ClientError, asyncio.TimeoutError) as e:
             print(f"Error in chat request: {e}")
-            await get_credential()  # Refresh credential on error
+            if retries < MAX_RETRIES:
+                await asyncio.sleep(RETRY_DELAY)
+                return await handle_chat_request(model, messages, retries + 1)
             return None
 
+@app.post(API_ENDPOINT)
+async def chat_endpoint(request: ChatRequest):
+    try:
+        model = MODEL_MAPPING.get(request.model, request.model or MODEL_NAME)
+        response_content = await handle_chat_request(model, request.messages)
+        
+        if response_content is None:
+            return Response(
+                content="Failed to get response from API",
+                status_code=500
+            )
+
+        return {
+            "choices": [{
+                "message": {
+                    "role": "assistant",
+                    "content": response_content.strip()
+                },
+                "finish_reason": "stop",
+                "index": 0
+            }],
+            "model": model,
+            "object": "chat.completion"
+        }
+
+    except Exception as e:
+        print(f"Error handling chat request: {e}")
+        return Response(content="Internal server error", status_code=500)
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=PORT)