Upload app.py

app.py

@@ -307,34 +307,39 @@ class MultiAuthManager:
     def __init__(self, credentials):
         self.auth_managers = [AuthManager(email, password) for email, password in credentials]
         self.current_index = 0
-        self._last_rotation = time.time()
-        self._rotation_interval = 300  # 5分钟轮转间隔
-
-    def _should_rotate(self) -> bool:
-        """检查是否应该轮转到下一个账号"""
-        return time.time() - self._last_rotation >= self._rotation_interval
+        self.last_success_index = 0  # 记录上一次成功的账号索引
+        self.last_rotation_date = datetime.now().date()  # 记录上次轮询的日期
 
     def get_next_auth_manager(self, model):
-        """改进的账号选择逻辑"""
-        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
+        """改进的账号选择逻辑，从上次成功的账号开始尝试"""
+        current_date = datetime.now().date()
+        
+        # 如果是新的一天，重置为从第一个账号开始
+        if 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
+        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
 
     def ensure_valid_token(self, model):
-        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
+        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
         return None
 
     def reset_all_model_status(self):
@@ -445,7 +450,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
             }
         ]
@@ -472,22 +477,30 @@ def stream_notdiamond_response(response, model):
     buffer = ""
     full_content = ""
     
-    for chunk in response.iter_content(chunk_size=1024):
+    for chunk in response.iter_lines():
         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
+                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
                 
             except Exception as e:
                 logger.error(f"Error processing chunk: {e}")
@@ -496,20 +509,71 @@ 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_content(chunk_size=1024):
+        for chunk in response.iter_lines():
             if chunk:
-                content = chunk.decode('utf-8')
-                full_content += content
-                context_buffer.append(content)
+                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
         
         completion_tokens = count_tokens(full_content, model)
         total_tokens = prompt_tokens + completion_tokens
@@ -526,8 +590,7 @@ def handle_non_stream_response(response, model, prompt_tokens):
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": full_content,
-                        "context": ''.join(context_buffer)  # 包含完整上下文
+                        "content": full_content
                     },
                     "finish_reason": "stop"
                 }
@@ -545,24 +608,6 @@ 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:
@@ -862,54 +907,8 @@ def make_request(payload, auth_manager, 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()
-
+# 删除 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)