Update app.py

app.py

@@ -7,22 +7,25 @@ import uuid
 import re
 import socket
 from concurrent.futures import ThreadPoolExecutor
-from functools import lru_cache
-from cachetools import TTLCache
+from functools import lru_cache, wraps
+from typing import Dict, Any, Callable
 
 import requests
 import tiktoken
 from flask import Flask, Response, jsonify, request, stream_with_context
 from flask_cors import CORS
-from typing import Dict, Any
 from requests.adapters import HTTPAdapter
 from urllib3.util.connection import create_connection
 import urllib3
+from cachetools import TTLCache
 
 # Constants
 CHAT_COMPLETION_CHUNK = 'chat.completion.chunk'
 CHAT_COMPLETION = 'chat.completion'
 CONTENT_TYPE_EVENT_STREAM = 'text/event-stream'
+_BASE_URL = "https://chat.notdiamond.ai"
+_API_BASE_URL = "https://spuckhogycrxcbomznwo.supabase.co"
+_USER_AGENT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36'
 
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -72,134 +75,144 @@ def create_custom_session():
 
 class AuthManager:
     def __init__(self, email: str, password: str):
-        self.email = email
-        self.password = password
-        self.api_key: str = ""
-        self.user_info: Dict[str, Any] = {}
-        self.refresh_token: str = ""
+        self._email: str = email
+        self._password: str = password
+        self._api_key: str = ""
+        self._user_info: Dict[str, Any] = {}
+        self._refresh_token: str = ""
+        self._access_token: str = ""
+        self._token_expiry: float = 0
+        self._session: requests.Session = create_custom_session()
         
-        self.logger = logging.getLogger(__name__)
+        self._logger: logging.Logger = logging.getLogger(__name__)
         logging.basicConfig(level=logging.INFO)
         
-        self.session = create_custom_session()
-        self.initialize_auth()
-        self.fetch_apikey()
-        self.log_values()
-    
-    def initialize_auth(self):
-        """初始化认证，优先使用缓存的 refresh_token"""
-        cache_key = f"{self.email}|{self.password}"
-        cached_token = refresh_token_cache.get(cache_key)
-        
-        if cached_token:
-            self.refresh_token = cached_token
-            self.refresh_user_token()
-        else:
-            self.login()
-    
-    def log_values(self) -> None:
+    def login(self) -> bool:
+        """使用电子邮件和密码进行用户登录,并获取用户信息。"""
+        url = f"{_API_BASE_URL}/auth/v1/token?grant_type=password"
+        headers = self._get_headers(with_content_type=True)
+        data = {
+            "email": self._email,
+            "password": self._password,
+            "gotrue_meta_security": {}
+        }
+
+        try:
+            response = self._make_request('POST', url, headers=headers, json=data)
+            self._user_info = response.json()
+            self._refresh_token = self._user_info.get('refresh_token', '')
+            self._access_token = self._user_info.get('access_token', '')
+            self._token_expiry = time.time() + self._user_info.get('expires_in', 3600)
+            self._log_values()
+            return True
+        except requests.RequestException as e:
+            self._logger.error(f"\033[91m登录请求错误: {e}\033[0m")
+            return False
+
+    def refresh_user_token(self) -> bool:
+        """使用刷新令牌来请求一个新的访问令牌并更新实例变量。"""
+        url = f"{_API_BASE_URL}/auth/v1/token?grant_type=refresh_token"
+        headers = self._get_headers(with_content_type=True)
+        data = {"refresh_token": self._refresh_token}
+
+        try:
+            response = self._make_request('POST', url, headers=headers, json=data)
+            self._user_info = response.json()
+            self._refresh_token = self._user_info.get('refresh_token', '')
+            self._access_token = self._user_info.get('access_token', '')
+            self._token_expiry = time.time() + self._user_info.get('expires_in', 3600)
+            self._log_values()
+            return True
+        except requests.RequestException as e:
+            self._logger.error(f"刷新令牌请求错误: {e}")
+            return False
+
+    def get_jwt_value(self) -> str:
+        """返回访问令牌。"""
+        return self._access_token
+
+    def is_token_valid(self) -> bool:
+        """检查当前的访问令牌是否有效。"""
+        return bool(self._access_token) and time.time() < self._token_expiry
+
+    def ensure_valid_token(self) -> bool:
+        """确保有一个有效的访问令牌，如果需要则刷新或重新登录。"""
+        if self.is_token_valid():
+            return True
+        if self._refresh_token:
+            if self.refresh_user_token():
+                return True
+        return self.login()
+
+    def clear_auth(self) -> None:
+        """清除当前的授权信息。"""
+        self._user_info = {}
+        self._refresh_token = ""
+        self._access_token = ""
+        self._token_expiry = 0
+
+    def _log_values(self) -> None:
         """记录刷新令牌到日志中。"""
-        self.logger.info(f"\033[92mRefresh Token: {self.refresh_token}\033[0m")
+        self._logger.info(f"\033[92mRefresh Token: {self._refresh_token}\033[0m")
+        self._logger.info(f"\033[92mAccess Token: {self._access_token}\033[0m")
         
-    def fetch_apikey(self) -> str:
+    def _fetch_apikey(self) -> str:
         """获取API密钥。"""
-        if self.api_key:
-            return self.api_key
-    
+        if self._api_key:
+            return self._api_key
+
         try:
-            url = "https://chat.notdiamond.ai/login"
-            headers = {
-                'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36'
-            }
-            response = self.session.get(url, headers=headers)
-            response.raise_for_status()
-    
-            # 匹配 <script> 标签中的 JS 文件路径
+            login_url = f"{_BASE_URL}/login"
+            response = self._make_request('GET', login_url)
+            
             match = re.search(r'<script src="(/_next/static/chunks/app/layout-[^"]+\.js)"', response.text)
             if not match:
-                self.logger.warning("未找到匹配的脚本标签")
-                return ""
-    
-            js_url = f"https://chat.notdiamond.ai{match.group(1)}"
-            js_response = self.session.get(js_url, headers=headers)
-            js_response.raise_for_status()
-            
-            # 匹配 API key
-            api_key_match = re.search(r'\("https://spuckhogycrxcbomznwo\.supabase\.co","([^"]+)"\)', js_response.text)
-            if api_key_match:
-                self.api_key = api_key_match.group(1)
-                logger.info(f"Extracted API key: {self.api_key}")
-                return self.api_key
-            else:
-                self.logger.error("未能匹配API key")
-                return ""
-    
-        except requests.RequestException as e:
-            self.logger.error(f"请求JS文件时发生错误: {e}")
-            return ""
+                raise ValueError("未找到匹配的脚本标签")
 
-    def login(self) -> None:
-        """使用电子邮件和密码进行用户登录，并获取用户信息。"""
-        api_key = self.fetch_apikey()
-        if not api_key:
-            self.logger.error("API key is missing, cannot proceed with login.")
-            return
-    
-        url = "https://spuckhogycrxcbomznwo.supabase.co/auth/v1/token?grant_type=password"
-        headers = {
-            'apikey': api_key,
-            'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
-            'Content-Type': 'application/json'
-        }
-        data = {
-            "email": self.email,
-            "password": self.password,
-            "gotrue_meta_security": {}
-        }
-    
-        try:
-            response = self.session.post(url, headers=headers, json=data)
-            response.raise_for_status()
-            self.user_info = response.json()
-            self.refresh_token = self.user_info.get('refresh_token', '')
+            js_url = f"{_BASE_URL}{match.group(1)}"
+            js_response = self._make_request('GET', js_url)
             
-            # 缓存 refresh_token
-            cache_key = f"{self.email}|{self.password}"
-            refresh_token_cache[cache_key] = self.refresh_token
+            api_key_match = re.search(r'$"https://spuckhogycrxcbomznwo\.supabase\.co","([^"]+)"$', js_response.text)
+            if not api_key_match:
+                raise ValueError("未能匹配API key")
             
-            logger.info(f"Login successful for email: {self.email}")
-    
-        except requests.RequestException as e:
-            self.logger.error(f"登录请求错误: {e}")
+            self._api_key = api_key_match.group(1)
+            return self._api_key
 
-    def refresh_user_token(self) -> None:
-        """使用刷新令牌来请求一个新的访问令牌并更新实例变量。"""
-        url = "https://spuckhogycrxcbomznwo.supabase.co/auth/v1/token?grant_type=refresh_token"
+        except (requests.RequestException, ValueError) as e:
+            self._logger.error(f"获取API密钥时发生错误: {e}")
+            return ""
+
+    def _get_headers(self, with_content_type: bool = False) -> Dict[str, str]:
+        """生成请求头。"""
         headers = {
-            'apikey': self.fetch_apikey(),
-            'content-type': 'application/json;charset=UTF-8',
-            'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36'
+            'apikey': self._fetch_apikey(),
+            'user-agent': _USER_AGENT
         }
-        data = {"refresh_token": self.refresh_token}
-
+        if with_content_type:
+            headers['Content-Type'] = 'application/json'
+        if self._access_token:
+            headers['Authorization'] = f'Bearer {self._access_token}'
+        return headers
+
+    def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
+        """发送HTTP请求并处理异常。"""
         try:
-            response = self.session.post(url, headers=headers, json=data)
+            response = self._session.request(method, url, **kwargs)
             response.raise_for_status()
-            self.user_info = response.json()
-            self.refresh_token = self.user_info.get('refresh_token', '')
-            
-            # 更新缓存中的 refresh_token
-            cache_key = f"{self.email}|{self.password}"
-            refresh_token_cache[cache_key] = self.refresh_token
-            
+            return response
         except requests.RequestException as e:
-            self.logger.error(f"刷新令牌请求错误: {e}")
-            # 如果刷新失败，尝试重新登录
-            self.login()
-
-    def get_jwt_value(self) -> str:
-        """返回访问令牌。"""
-        return self.user_info.get('access_token', '')
+            self._logger.error(f"请求错误 ({method} {url}): {e}")
+            raise
+
+def require_auth(func: Callable) -> Callable:
+    """装饰器，确保在调用API之前有有效的token。"""
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not self.ensure_valid_token():
+            raise Exception("无法获取有效的授权token")
+        return func(self, *args, **kwargs)
+    return wrapper
 
 # 全局的 AuthManager 对象，将在每次请求时更新
 auth_manager = None
@@ -217,9 +230,7 @@ def get_notdiamond_headers():
         'accept': 'text/event-stream',
         'accept-language': 'zh-CN,zh;q=0.9',
         'content-type': 'application/json',
-        'user-agent': ('Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) '
-                       'AppleWebKit/537.36 (KHTML, like Gecko) '
-                       'Chrome/128.0.0.0 Safari/537.36'),
+        'user-agent': _USER_AGENT,
         'authorization': f'Bearer {auth_manager.get_jwt_value()}'
     }
 
@@ -334,7 +345,7 @@ def handle_non_stream_response(response, model, prompt_tokens):
         "object": "chat.completion",
         "created": int(time.time()),
         "model": model,
-        "system_fingerprint": generate_system_fingerprint(),
+                "system_fingerprint": generate_system_fingerprint(),
         "choices": [
             {
                 "index": 0,
@@ -359,6 +370,7 @@ def generate_stream_response(response, model, prompt_tokens):
     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,
@@ -372,12 +384,8 @@ def generate_stream_response(response, model, prompt_tokens):
 def get_auth_credentials():
     """从请求头中获取认证凭据"""
     auth_header = request.headers.get('Authorization')
-    if not auth_header:
-        logger.error("Authorization header is missing")
-        return None, None
-    
-    if not auth_header.startswith('Bearer '):
-        logger.error(f"Authorization header format is incorrect: {auth_header}")
+    if not auth_header or not auth_header.startswith('Bearer '):
+        logger.error("Authorization header is missing or invalid")
         return None, None
     
     try:
@@ -395,6 +403,8 @@ def before_request():
     email, password = get_auth_credentials()
     if email and password:
         auth_manager = AuthManager(email, password)
+        if not auth_manager.ensure_valid_token():
+            auth_manager = None
     else:
         auth_manager = None
 
@@ -542,6 +552,7 @@ def build_payload(request_data, model_id):
 
 def make_request(payload):
     """发送请求并处理可能的认证刷新。"""
+    global auth_manager
     url = get_notdiamond_url()
     headers = get_notdiamond_headers()
     response = executor.submit(requests.post, url, headers=headers, json=payload, stream=True).result()
@@ -562,3 +573,4 @@ def make_request(payload):
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 3000))
     app.run(debug=False, host='0.0.0.0', port=port, threaded=True)
+