Upload 2 files

Dockerfile

@@ -23,5 +23,4 @@ ENV PYTHONUNBUFFERED=1
 EXPOSE 3000
 
 # 使用 gunicorn 作为生产级 WSGI 服务器
-# 在 Dockerfile 中修改 gunicorn 命令
-CMD ["gunicorn", "--bind", "0.0.0.0:3000", "--workers", "4", "--timeout", "120", "--keep-alive", "5", "--worker-class", "sync", "app:app"]
+CMD ["gunicorn", "--bind", "0.0.0.0:3000", "--workers", "4", "app:app"]

app.py

@@ -18,6 +18,8 @@ from urllib3.util.connection import create_connection
 import urllib3
 from cachetools import TTLCache
 import threading
+from time import sleep
+from datetime import datetime, timedelta
 
 # 新增导入
 import register_bot
@@ -98,14 +100,8 @@ class CustomHTTPAdapter(HTTPAdapter):
 
 # 创建自定义的 Session
 def create_custom_session():
-    """创建自定义的 Session，添加超时设置"""
     session = requests.Session()
-    adapter = CustomHTTPAdapter(
-        pool_connections=100,
-        pool_maxsize=100,
-        max_retries=3,
-        pool_block=False
-    )
+    adapter = CustomHTTPAdapter()
     session.mount('https://', adapter)
     session.mount('http://', adapter)
     return session
@@ -311,42 +307,18 @@ class MultiAuthManager:
     def __init__(self, credentials):
         self.auth_managers = [AuthManager(email, password) for email, password in credentials]
         self.current_index = 0
-        self.last_success_index = None
         self._last_rotation = time.time()
-        self._rotation_interval = 300
-        self._model_usage = {}
-        self._invalid_accounts = set()  # 记录失效的账号
-
-    def remove_invalid_account(self, auth_manager):
-        """移除失效的账号"""
-        if auth_manager._email not in self._invalid_accounts:
-            self._invalid_accounts.add(auth_manager._email)
-            self.auth_managers = [am for am in self.auth_managers if am._email != auth_manager._email]
-            logger.info(f"已移除失效账号: {auth_manager._email}")
-            
-            # 如果移除的是当前使用的账号，重置索引
-            if self.last_success_index is not None:
-                if self.last_success_index >= len(self.auth_managers):
-                    self.last_success_index = None
-            if self.current_index >= len(self.auth_managers):
-                self.current_index = 0
+        self._rotation_interval = 300  # 5分钟轮转间隔
+
+    def _should_rotate(self) -> bool:
+        """检查是否应该轮转到下一个账号"""
+        return time.time() - self._last_rotation >= self._rotation_interval
 
     def get_next_auth_manager(self, model):
-        """改进的账号选择逻辑，优先使用模型对应的上次成功账号"""
-        if not self.auth_managers:  # 如果没有可用账号，返回 None
-            return None
-
-        # 首先尝试使用该模型上次成功的账号
-        if model in self._model_usage:
-            last_success_index = self._model_usage[model]
-            if last_success_index < len(self.auth_managers):
-                auth_manager = self.auth_managers[last_success_index]
-                if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
-                    return auth_manager
-        
-        # 如果没有该模型的成功记录，或上次成功的账号不可用，则从当前位置开始轮询
-        if len(self.auth_managers) == 0:
-            return None
+        """改进的账号选择逻辑"""
+        if self._should_rotate():
+            self.current_index = (self.current_index + 1) % len(self.auth_managers)
+            self._last_rotation = time.time()
 
         start_index = self.current_index
         for _ in range(len(self.auth_managers)):
@@ -359,34 +331,13 @@ class MultiAuthManager:
         return None
 
     def ensure_valid_token(self, model):
-        """确保获取有效的token并返回可用的auth_manager"""
-        auth_manager = self.get_next_auth_manager(model)
-        if not auth_manager:
-            return None
-            
-        try:
-            if auth_manager.ensure_valid_token():
+        for _ in range(len(self.auth_managers)):
+            auth_manager = self.get_next_auth_manager(model)
+            if auth_manager and auth_manager.ensure_valid_token():
                 return auth_manager
-        except requests.exceptions.RequestException as e:
-            if "400 Client Error: Bad Request" in str(e):
-                logger.error(f"账号 {auth_manager._email} 已失效")
-                self.remove_invalid_account(auth_manager)
-            return None
-            
         return None
 
-    def mark_success(self, auth_manager, model):
-        """记录成功使用的账号索引，并与模型关联"""
-        for i, manager in enumerate(self.auth_managers):
-            if manager == auth_manager:
-                self._model_usage[model] = i
-                self.last_success_index = i
-                break
-
-    def reset_model_status(self):
-        """重置所有账号的模型使用状态"""
-        self._model_usage.clear()
-        self._invalid_accounts.clear()  # 清除失效账号记录
+    def reset_all_model_status(self):
         for auth_manager in self.auth_managers:
             auth_manager.reset_model_status()
 
@@ -516,54 +467,37 @@ def count_message_tokens(messages, model="gpt-3.5-turbo-0301"):
     """计算消息列表中的总令牌数量。"""
     return sum(count_tokens(str(message), model) for message in messages)
 
-# 在文件开头添加常量
-STREAM_TIMEOUT = 30  # 流式响应超时时间（秒）
-REQUEST_TIMEOUT = 10  # 普通请求超时时间（秒）
-CHUNK_SIZE = 512  # 减小块大小以加快处理速度
-
 def stream_notdiamond_response(response, model):
-    """改进的流式响应处理，添加超时和错误处理"""
+    """改进的流式响应处理，确保保持上下文完整性。"""
     buffer = ""
     full_content = ""
-    last_chunk_time = time.time()
     
-    try:
-        for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
-            current_time = time.time()
-            if current_time - last_chunk_time > STREAM_TIMEOUT:
-                logger.warning("Stream timeout reached")
-                break
+    for chunk in response.iter_content(chunk_size=1024):
+        if chunk:
+            try:
+                new_content = chunk.decode('utf-8')
+                buffer += new_content
+                full_content += new_content
                 
-            if chunk:
-                try:
-                    new_content = chunk.decode('utf-8')
-                    buffer += new_content
-                    full_content += new_content
-                    
-                    # 创建完整的响应块
-                    chunk_data = create_openai_chunk(new_content, model)
-                    
-                    # 确保响应块包含完整的上下文
-                    if 'choices' in chunk_data and chunk_data['choices']:
-                        chunk_data['choices'][0]['delta']['content'] = new_content
-                        chunk_data['choices'][0]['context'] = full_content
-                    
-                    yield chunk_data
-                    last_chunk_time = current_time
-                    
-                except Exception as e:
-                    logger.error(f"Error processing chunk: {e}")
-                    continue
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Stream error: {e}")
-    except Exception as e:
-        logger.error(f"Unexpected error in stream processing: {e}")
-    finally:
-        # 确保发送完成标记
-        final_chunk = create_openai_chunk('', model, 'stop')
-        if 'choices' in final_chunk and final_chunk['choices']:
-            final_chunk['choices'][0]['context'] = full_content
-        yield final_chunk
+                # 创建完整的响应块
+                chunk_data = create_openai_chunk(new_content, model)
+                
+                # 确保响应块包含完整的上下文
+                if 'choices' in chunk_data and chunk_data['choices']:
+                    chunk_data['choices'][0]['delta']['content'] = new_content
+                    chunk_data['choices'][0]['context'] = full_content  # 添加完整上下文
+                
+                yield chunk_data
+                
+            except Exception as e:
+                logger.error(f"Error processing chunk: {e}")
+                continue
+    
+    # 发送完成标记
+    final_chunk = create_openai_chunk('', model, 'stop')
+    if 'choices' in final_chunk and final_chunk['choices']:
+        final_chunk['choices'][0]['context'] = full_content  # 在最终块中包含完整上下文
+    yield final_chunk
 
 def handle_non_stream_response(response, model, prompt_tokens):
     """改进的非流式响应处理，确保保持完整上下文。"""
@@ -611,32 +545,23 @@ def handle_non_stream_response(response, model, prompt_tokens):
         logger.error(f"Error processing non-stream response: {e}")
         raise
 
-# 修改 generate_stream_response 函数
 def generate_stream_response(response, model, prompt_tokens):
-    """改进的流式 HTTP 响应生成器"""
+    """生成流式 HTTP 响应。"""
     total_completion_tokens = 0
-    start_time = time.time()
     
-    try:
-        for chunk in stream_notdiamond_response(response, model):
-            if time.time() - start_time > STREAM_TIMEOUT:
-                logger.warning("Response generation timeout")
-                break
-                
-            content = chunk['choices'][0]['delta'].get('content', '')
-            total_completion_tokens += count_tokens(content, model)
-            
-            chunk['usage'] = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": total_completion_tokens,
-                "total_tokens": prompt_tokens + total_completion_tokens
-            }
-            
-            yield f"data: {json.dumps(chunk)}\n\n"
-    except Exception as e:
-        logger.error(f"Error generating stream response: {e}")
-    finally:
-        yield "data: [DONE]\n\n"
+    for chunk in stream_notdiamond_response(response, model):
+        content = chunk['choices'][0]['delta'].get('content', '')
+        total_completion_tokens += count_tokens(content, model)
+        
+        chunk['usage'] = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": total_completion_tokens,
+            "total_tokens": prompt_tokens + total_completion_tokens
+        }
+        
+        yield f"data: {json.dumps(chunk)}\n\n"
+    
+    yield "data: [DONE]\n\n"
 
 def get_auth_credentials():
     """从API获取认证凭据"""
@@ -898,57 +823,93 @@ def make_request(payload, auth_manager, model_id):
                     url, 
                     headers=headers, 
                     json=payload, 
-                    stream=True,
-                    timeout=REQUEST_TIMEOUT  # 添加超时设置
+                    stream=True
                 ).result()
                 
                 if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
                     logger.info(f"请求成功，使用账号 {auth_manager._email}")
-                    multi_auth_manager.mark_success(auth_manager, model_id)
                     return response
                 
                 headers_cache.clear()
                 
                 if response.status_code == 401:  # Unauthorized
                     logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
-                    try:
-                        if auth_manager.ensure_valid_token():
-                            continue
-                    except requests.exceptions.RequestException as e:
-                        if "400 Client Error: Bad Request" in str(e):
-                            logger.error(f"账号 {auth_manager._email} 已失效")
-                            multi_auth_manager.remove_invalid_account(auth_manager)
-                            break
+                    if auth_manager.ensure_valid_token():
+                        continue
                 
                 if response.status_code == 403:  # Forbidden, 模型使用限制
                     logger.warning(f"Model {model_id} usage limit reached for account {auth_manager._email}")
                     auth_manager.set_model_unavailable(model_id)
-                    break
+                    break  # 跳出重试循环，尝试下一个账号
                 
                 logger.error(f"Request failed with status {response.status_code} for account {auth_manager._email}")
                 
             except Exception as e:
                 logger.error(f"Request attempt {attempt + 1} failed for account {auth_manager._email}: {e}")
-                if "400 Client Error: Bad Request" in str(e):
-                    logger.error(f"账号 {auth_manager._email} 已失效")
-                    multi_auth_manager.remove_invalid_account(auth_manager)
-                    break
                 if attempt < max_retries - 1:
                     time.sleep(retry_delay)
                 continue
 
-    # 检查是否需要注册新账号
-    if not multi_auth_manager.auth_managers or len(tried_accounts) == len(multi_auth_manager.auth_managers):
-        logger.info("所有现有账号都已尝试或无可用账号，开始注册新账号")
+    # 所有账号都尝试过且失败后，才进行注册
+    if len(tried_accounts) == len(multi_auth_manager.auth_managers):
+        logger.info("所有现有账号都已尝试，开始注册新账号")
         successful_accounts = register_bot.register_and_verify(5)
         if successful_accounts:
             credentials = [(account['email'], account['password']) for account in successful_accounts]
             multi_auth_manager = MultiAuthManager(credentials)
+            # 使用新注册的账号重试请求
             return make_request(payload, None, model_id)
     
     raise Exception("所有账号均不可用，且注册新账号失败")
 
+def health_check():
+    """改进的健康检查函数"""
+    last_check_time = {}  # 用于跟踪每个账号的最后检查时间
+    
+    while True:
+        try:
+            if multi_auth_manager:
+                current_time = time.time()
+                
+                for auth_manager in multi_auth_manager.auth_managers:
+                    email = auth_manager._email
+                    
+                    # 检查是否需要进行健康检查
+                    if email not in last_check_time or \
+                       current_time - last_check_time[email] >= AUTH_CHECK_INTERVAL:
+                        
+                        if not auth_manager._should_attempt_auth():
+                            logger.info(f"Skipping health check for {email} due to rate limiting")
+                            continue
+                            
+                        if not auth_manager.ensure_valid_token():
+                            logger.warning(f"Auth token validation failed during health check for {email}")
+                            auth_manager.clear_auth()
+                        else:
+                            logger.info(f"Health check passed for {email}")
+                            
+                        last_check_time[email] = current_time
+                
+                # 每天重置所有账号的模型使用状态
+                current_time_local = time.localtime()
+                if current_time_local.tm_hour == 0 and current_time_local.tm_min == 0:
+                    multi_auth_manager.reset_all_model_status()
+                    logger.info("Reset model status for all accounts")
+                    
+        except Exception as e:
+            logger.error(f"Health check error: {e}")
+            
+        sleep(60)  # 主循环每分钟运行一次
+
+# 为了兼容 Flask CLI 和 Gunicorn，修改启动逻辑
+if __name__ != "__main__":
+    health_check_thread = threading.Thread(target=health_check, daemon=True)
+    health_check_thread.start()
+
 if __name__ == "__main__":
+    health_check_thread = threading.Thread(target=health_check, daemon=True)
+    health_check_thread.start()
+    
     port = int(os.environ.get("PORT", 3000))
     app.run(debug=False, host='0.0.0.0', port=port, threaded=True)