Upload 2 files

Dockerfile

@@ -23,4 +23,4 @@ ENV PYTHONUNBUFFERED=1
 EXPOSE 3000
 
 # 使用 gunicorn 作为生产级 WSGI 服务器
-CMD ["gunicorn", "--bind", "0.0.0.0:3000", "--workers", "4", "--timeout", "120", "app:app"]
+CMD ["gunicorn", "--bind", "0.0.0.0:3000", "--workers", "4", "app:app"]

app.py

@@ -107,7 +107,7 @@ def create_custom_session():
     return session
 
 # 添加速率限制相关的常量
-AUTH_RETRY_DELAY = 60  # 认证试延迟（秒）
+AUTH_RETRY_DELAY = 60  # 认证重试延迟（秒）
 AUTH_BACKOFF_FACTOR = 2  # 退避因子
 AUTH_MAX_RETRIES = 3  # 最大重试次数
 AUTH_CHECK_INTERVAL = 300  # 健康检查间隔（秒）
@@ -307,42 +307,34 @@ 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 = 0  # 记录上一次成功的账号索引
-        self._last_rotation_date = None  # 使用None作为初始值
-        self._lock = threading.Lock()  # 添加线程锁
+        self._last_rotation = time.time()
+        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):
-        """改进的账号选择逻辑，从上次成功的账号开始尝试"""
-        with self._lock:  # 使用线程锁保护共享状态
-            try:
-                current_date = datetime.now().date()
-                
-                # 初始化或检查日期变更
-                if self._last_rotation_date is None or current_date > self._last_rotation_date:
-                    self.current_index = 0
-                    self.last_success_index = 0
-                    self._last_rotation_date = current_date
-                    # 重置所有账号的模型状态
-                    for auth_manager in self.auth_managers:
-                        auth_manager.reset_model_status()
-                    return self.auth_managers[0] if self.auth_managers else None
-                    
-                # 从上次成功的账号开始尝试
-                self.current_index = self.last_success_index
-                if 0 <= self.current_index < len(self.auth_managers):
-                    auth_manager = self.auth_managers[self.current_index]
-                    if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
-                        return auth_manager
-                return None
-            except Exception as e:
-                logger.error(f"Error in get_next_auth_manager: {str(e)}")
-                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)):
+            auth_manager = self.auth_managers[self.current_index]
+            if auth_manager.is_model_available(model) and auth_manager._should_attempt_auth():
+                return auth_manager
+            self.current_index = (self.current_index + 1) % len(self.auth_managers)
+            if self.current_index == start_index:
+                break
+        return None
 
     def ensure_valid_token(self, model):
-        auth_manager = self.get_next_auth_manager(model)
-        if auth_manager and auth_manager.ensure_valid_token():
-            self.last_success_index = self.current_index  # 更新最后成功的账号索引
-            return auth_manager
+        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
         return None
 
     def reset_all_model_status(self):
@@ -453,7 +445,7 @@ def create_openai_chunk(content, model, finish_reason=None, usage=None):
                 "delta": {"content": content} if content else {},
                 "logprobs": None,
                 "finish_reason": finish_reason,
-                # 添加上下文关信息
+                # 添加上下文相关信息
                 "context_preserved": True
             }
         ]
@@ -480,30 +472,22 @@ def stream_notdiamond_response(response, model):
     buffer = ""
     full_content = ""
     
-    for chunk in response.iter_lines():
+    for chunk in response.iter_content(chunk_size=1024):
         if chunk:
             try:
-                chunk_str = chunk.decode('utf-8')
-                # 跳过SSE前缀
-                if chunk_str.startswith('data: '):
-                    chunk_str = chunk_str[6:]
-                elif chunk_str == 'data: [DONE]':
-                    continue
-                    
-                # 尝试解析JSON
-                try:
-                    chunk_data = json.loads(chunk_str)
-                    content = chunk_data.get('choices', [{}])[0].get('delta', {}).get('content', '')
-                    if content:
-                        full_content += content
-                        chunk_data = create_openai_chunk(content, model)
-                        yield chunk_data
-                except json.JSONDecodeError:
-                    # 如果不是JSON格式，直接作为内容处理
-                    if chunk_str.strip():
-                        full_content += chunk_str
-                        chunk_data = create_openai_chunk(chunk_str, model)
-                        yield chunk_data
+                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
                 
             except Exception as e:
                 logger.error(f"Error processing chunk: {e}")
@@ -512,71 +496,20 @@ def stream_notdiamond_response(response, model):
     # 发送完成标记
     final_chunk = create_openai_chunk('', model, 'stop')
     if 'choices' in final_chunk and final_chunk['choices']:
-        final_chunk['choices'][0]['context'] = full_content
+        final_chunk['choices'][0]['context'] = full_content  # 在最终块中包含完整上下文
     yield final_chunk
 
-def generate_stream_response(response, model, prompt_tokens):
-    """生成流式 HTTP 响应，确保完整的上下文。"""
-    total_completion_tokens = 0
-    full_content = ""
-    
-    for chunk in stream_notdiamond_response(response, model):
-        content = chunk['choices'][0]['delta'].get('content', '')
-        if content:
-            full_content += content
-            total_completion_tokens = count_tokens(full_content, model)
-            
-            chunk['usage'] = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": total_completion_tokens,
-                "total_tokens": prompt_tokens + total_completion_tokens
-            }
-            
-            # 确保每个块都包含完整的上下文
-            chunk['choices'][0]['context'] = full_content
-            
-            yield f"data: {json.dumps(chunk)}\n\n"
-    
-    # 发送最终的完成标记
-    final_chunk = create_openai_chunk('', model, 'stop')
-    final_chunk['choices'][0]['context'] = full_content
-    final_chunk['usage'] = {
-        "prompt_tokens": prompt_tokens,
-        "completion_tokens": total_completion_tokens,
-        "total_tokens": prompt_tokens + total_completion_tokens
-    }
-    yield f"data: {json.dumps(final_chunk)}\n\n"
-    yield "data: [DONE]\n\n"
-
 def handle_non_stream_response(response, model, prompt_tokens):
-    """改进的非流式响应处理，确保完整的上下文。"""
+    """改进的非流式响应处理，确保保持完整上下文。"""
     full_content = ""
+    context_buffer = []
     
     try:
-        for chunk in response.iter_lines():
+        for chunk in response.iter_content(chunk_size=1024):
             if chunk:
-                try:
-                    chunk_str = chunk.decode('utf-8')
-                    # 跳过SSE前缀
-                    if chunk_str.startswith('data: '):
-                        chunk_str = chunk_str[6:]
-                    elif chunk_str == 'data: [DONE]':
-                        continue
-                        
-                    # 尝试解析JSON
-                    try:
-                        chunk_data = json.loads(chunk_str)
-                        content = chunk_data.get('choices', [{}])[0].get('delta', {}).get('content', '')
-                        if content:
-                            full_content += content
-                    except json.JSONDecodeError:
-                        # 如果不是JSON格式，直接作为内容处理
-                        if chunk_str.strip():
-                            full_content += chunk_str
-                            
-                except Exception as e:
-                    logger.error(f"Error processing chunk in non-stream response: {e}")
-                    continue
+                content = chunk.decode('utf-8')
+                full_content += content
+                context_buffer.append(content)
         
         completion_tokens = count_tokens(full_content, model)
         total_tokens = prompt_tokens + completion_tokens
@@ -593,7 +526,8 @@ def handle_non_stream_response(response, model, prompt_tokens):
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": full_content
+                        "content": full_content,
+                        "context": ''.join(context_buffer)  # 包含完整上下文
                     },
                     "finish_reason": "stop"
                 }
@@ -611,6 +545,24 @@ def handle_non_stream_response(response, model, prompt_tokens):
         logger.error(f"Error processing non-stream response: {e}")
         raise
 
+def generate_stream_response(response, model, prompt_tokens):
+    """生成流式 HTTP 响应。"""
+    total_completion_tokens = 0
+    
+    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获取认证凭据"""
     try:
@@ -831,7 +783,6 @@ def make_request(payload, auth_manager, model_id):
     global multi_auth_manager
     max_retries = 3
     retry_delay = 1
-    request_timeout = 10  # 减少单次请求超时时间
     
     logger.info(f"尝试发送请求，模型：{model_id}")
     
@@ -851,92 +802,114 @@ def make_request(payload, auth_manager, model_id):
     # 记录已尝试的账号
     tried_accounts = set()
     
-    while True:  # 持续尝试，直到成功或确定所有账号都不可用
-        try:
-            auth_manager = multi_auth_manager.get_next_auth_manager(model_id)
-            if not auth_manager:
-                # 检查是否所有账号都已尝试
-                if len(tried_accounts) >= len(multi_auth_manager.auth_managers):
-                    logger.error("所有账号都已尝试且不可用")
-                    # 尝试注册新账号
-                    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)
-                        tried_accounts.clear()  # 清空已尝试账号列表
-                        continue  # 继续尝试新注册的账号
-                    else:
-                        raise Exception("所有账号均不可用，且注册新账号失败")
-                continue  # 如果还有未尝试的账号，继续循环
+    while len(tried_accounts) < len(multi_auth_manager.auth_managers):
+        auth_manager = multi_auth_manager.get_next_auth_manager(model_id)
+        if not auth_manager:
+            break
+            
+        # 如果这个账号已经尝试过，继续下一个
+        if auth_manager._email in tried_accounts:
+            continue
+            
+        tried_accounts.add(auth_manager._email)
+        logger.info(f"尝试使用账号 {auth_manager._email}")
+
+        for attempt in range(max_retries):
+            try:
+                url = get_notdiamond_url()
+                headers = get_notdiamond_headers(auth_manager)
+                response = executor.submit(
+                    requests.post, 
+                    url, 
+                    headers=headers, 
+                    json=payload, 
+                    stream=True
+                ).result()
                 
-            # 如果这个账号已经尝试过，跳过
-            if auth_manager._email in tried_accounts:
-                continue
+                if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
+                    logger.info(f"请求成功，使用账号 {auth_manager._email}")
+                    return response
+                
+                headers_cache.clear()
+                
+                if response.status_code == 401:  # Unauthorized
+                    logger.info(f"Token expired for account {auth_manager._email}, attempting refresh")
+                    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  # 跳出重试循环，尝试下一个账号
                 
-            tried_accounts.add(auth_manager._email)
-            logger.info(f"尝试使用账号 {auth_manager._email}")
+                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 attempt < max_retries - 1:
+                    time.sleep(retry_delay)
+                continue
 
-            for attempt in range(max_retries):
-                try:
-                    url = get_notdiamond_url()
-                    headers = get_notdiamond_headers(auth_manager)
-                    
-                    # 使��� session 发送请求
-                    session = create_custom_session()
-                    response = session.post(
-                        url,
-                        headers=headers,
-                        json=payload,
-                        stream=True,
-                        timeout=request_timeout
-                    )
-                    
-                    if response.status_code == 200:
-                        logger.info(f"请求成功，使用账号 {auth_manager._email}")
-                        return response
-                    
-                    # 处理不同的错误状态码
-                    if response.status_code == 401:  # Unauthorized
-                        logger.warning(f"Token expired for account {auth_manager._email}")
-                        if auth_manager.ensure_valid_token():
-                            continue
-                        break
-                    
-                    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
+    # 所有账号都尝试过且失败后，才进行注册
+    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 response.status_code >= 500:  # Server error
-                        if attempt < max_retries - 1:
-                            time.sleep(retry_delay * (attempt + 1))
+                    # 检查是否需要进行健康检查
+                    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")
                     
-                    logger.error(f"Request failed with status {response.status_code}")
-                    
-                except requests.Timeout:
-                    logger.warning(f"Request timeout for account {auth_manager._email}, attempt {attempt + 1}")
-                    if attempt < max_retries - 1:
-                        time.sleep(retry_delay * (attempt + 1))
-                    continue
-                except requests.RequestException as e:
-                    logger.error(f"Request error: {str(e)}")
-                    if attempt < max_retries - 1:
-                        time.sleep(retry_delay * (attempt + 1))
-                    continue
-                except Exception as e:
-                    logger.error(f"Unexpected error: {str(e)}")
-                    if attempt < max_retries - 1:
-                        time.sleep(retry_delay * (attempt + 1))
-                    continue
         except Exception as e:
-            logger.error(f"Error in main request loop: {str(e)}")
-            time.sleep(retry_delay)
-    
-    raise Exception("无法完成请求，所有可用账号都已尝试")
+            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()
 
-# 删除 health_check 函数和相关的线程启动代码
 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)