Update app.py

app.py

@@ -8,6 +8,7 @@ import re
 import socket
 from concurrent.futures import ThreadPoolExecutor
 from functools import lru_cache
+from cachetools import TTLCache
 
 import requests
 import tiktoken
@@ -40,6 +41,9 @@ if not NOTDIAMOND_IP:
     logger.error("NOTDIAMOND_IP environment variable is not set!")
     raise ValueError("NOTDIAMOND_IP must be set")
 
+# 创建一个 TTLCache 来存储 refresh_token
+refresh_token_cache = TTLCache(maxsize=1000, ttl=3600)
+
 # 自定义连接函数
 def patched_create_connection(address, *args, **kwargs):
     host, port = address
@@ -66,7 +70,6 @@ def create_custom_session():
     session.mount('http://', adapter)
     return session
 
-# 修改 AuthManager 类以使用自定义 Session
 class AuthManager:
     def __init__(self, email: str, password: str):
         self.email = email
@@ -79,10 +82,21 @@ class AuthManager:
         logging.basicConfig(level=logging.INFO)
         
         self.session = create_custom_session()
-        self.login()
+        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:
         """记录刷新令牌到日志中。"""
         self.logger.info(f"\033[92mRefresh Token: {self.refresh_token}\033[0m")
@@ -109,7 +123,7 @@ class AuthManager:
             js_response = self.session.get(js_url, headers=headers)
             js_response.raise_for_status()
             
-            api_key_match = re.search(r'\("https://spuckhogycrxcbomznwo\.supabase\.co","([^"]+)"\)', js_response.text)
+            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)
                 return self.api_key
@@ -140,6 +154,11 @@ class AuthManager:
             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
+            
         except requests.RequestException as e:
             self.logger.error(f"\033[91m登录请求错误: {e}\033[0m")
 
@@ -158,18 +177,22 @@ class AuthManager:
             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
+            
         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', '')
 
-# 初始化 AuthManager
-auth_manager = AuthManager(
-    os.getenv("AUTH_EMAIL", "[email protected]"),
-    os.getenv("AUTH_PASSWORD", "default_password"),
-)
+# 全局的 AuthManager 对象，将在每次请求时更新
+auth_manager = None
 
 NOTDIAMOND_URLS = os.getenv('NOTDIAMOND_URLS', 'https://not-diamond-workers.t7-cc4.workers.dev/stream-message').split(',')
 
@@ -263,7 +286,6 @@ def create_openai_chunk(content, model, finish_reason=None, usage=None):
         chunk["usage"] = usage
     return chunk
 
-
 def count_tokens(text, model="gpt-3.5-turbo-0301"):
     """计算给定文本的令牌数量。"""
     try:
@@ -327,7 +349,6 @@ 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,
@@ -338,6 +359,28 @@ def generate_stream_response(response, model, prompt_tokens):
     
     yield "data: [DONE]\n\n"
 
+def get_auth_credentials():
+    """从请求头中获取认证凭据"""
+    auth_header = request.headers.get('Authorization')
+    if not auth_header or not auth_header.startswith('Bearer '):
+        return None, None
+    
+    try:
+        credentials = auth_header.split('Bearer ')[1]
+        email, password = credentials.split('|')
+        return email.strip(), password.strip()
+    except:
+        return None, None
+
+@app.before_request
+def before_request():
+    global auth_manager
+    email, password = get_auth_credentials()
+    if email and password:
+        auth_manager = AuthManager(email, password)
+    else:
+        auth_manager = None
+
 @app.route('/', methods=['GET'])
 def root():
     return jsonify({
@@ -385,6 +428,10 @@ def proxy_models():
 @app.route('/ai/v1/chat/completions', methods=['POST'])
 def handle_request():
     """处理聊天完成请求。"""
+    global auth_manager
+    if not auth_manager:
+        return jsonify({'error': 'Unauthorized'}), 401
+
     try:
         request_data = request.get_json()
         model_id = request_data.get('model', '')
@@ -480,39 +527,20 @@ def make_request(payload):
     """发送请求并处理可能的认证刷新。"""
     url = get_notdiamond_url()
     headers = get_notdiamond_headers()
-    session = create_custom_session()
-    logger.info(f"Sending request to URL: {url}")
-    try:
-        response = executor.submit(session.post, url, headers=headers, json=payload, stream=True).result()
-        logger.info(f"Response status code: {response.status_code}")
-        if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
-            return response
-    except Exception as e:
-        logger.error(f"Error in make_request: {str(e)}")
-        raise
+    response = executor.submit(requests.post, url, headers=headers, json=payload, stream=True).result()
+    if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
+        return response
 
-    logger.info("Refreshing user token...")
     auth_manager.refresh_user_token()
     headers = get_notdiamond_headers()
-    try:
-        response = executor.submit(session.post, url, headers=headers, json=payload, stream=True).result()
-        logger.info(f"Response status code after token refresh: {response.status_code}")
-        if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
-            return response
-    except Exception as e:
-        logger.error(f"Error in make_request after token refresh: {str(e)}")
-        raise
+    response = executor.submit(requests.post, url, headers=headers, json=payload, stream=True).result()
+    if response.status_code == 200 and response.headers.get('Content-Type') == 'text/event-stream':
+        return response
 
-    logger.info("Logging in again...")
     auth_manager.login()
     headers = get_notdiamond_headers()
-    try:
-        response = executor.submit(session.post, url, headers=headers, json=payload, stream=True).result()
-        logger.info(f"Response status code after login: {response.status_code}")
-        return response
-    except Exception as e:
-        logger.error(f"Error in make_request after login: {str(e)}")
-        raise
+    response = executor.submit(requests.post, url, headers=headers, json=payload, stream=True).result()
+    return response
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 3000))