Upload 4 files

app.py

@@ -1,83 +1,104 @@
 from flask import Flask, request, jsonify, Response, stream_with_context, render_template_string
+from google.generativeai.types import generation_types
+from google.api_core.exceptions import InvalidArgument, ResourceExhausted, ServiceUnavailable, InternalServerError, Aborted
 import google.generativeai as genai
 import json
-from datetime import datetime
 import os
 import logging
 import func
+from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
-import requests
 import time
+import requests
+from collections import deque
+
 
 os.environ['TZ'] = 'Asia/Shanghai'
+
 app = Flask(__name__)
 
 app.secret_key = os.urandom(24)
 
-formatter = logging.Formatter('%(message)s') 
 
+formatter = logging.Formatter('%(message)s') 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO) 
-
+logger.setLevel(logging.INFO)
 handler = logging.StreamHandler()
 handler.setFormatter(formatter)
+logger.addHandler(handler)
 
-logger.addHandler(handler) 
-
-safety_settings = [
-    {
-        "category": "HARM_CATEGORY_HARASSMENT",
-        "threshold": "BLOCK_NONE"
-    },
-    {
-        "category": "HARM_CATEGORY_HATE_SPEECH",
-        "threshold": "BLOCK_NONE"
-    },
-    {
-        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-        "threshold": "BLOCK_NONE"
-    },
-    {
-        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-        "threshold": "BLOCK_NONE"
-    },
-]
+MAX_RETRIES = int(os.environ.get('MaxRetries', 3))
+MAX_REQUESTS = int(os.environ.get('MaxRequests', 4))
+LIMIT_WINDOW = int(os.environ.get('LimitWindow', 60))
+RETRY_DELAY = 1
+MAX_RETRY_DELAY = 16
+
+request_counts = {}
+
+api_key_blacklist = set()
+api_key_blacklist_duration = 60
 
 class APIKeyManager:
     def __init__(self):
         self.api_keys = os.environ.get('KeyArray').split(',')
-        self.current_index = 0  
-        logger.info(f"API Keys: {len(self.api_keys)} found")
+        self.current_index = 0
+
+    def get_available_key(self):
+        num_keys = len(self.api_keys)
+        for _ in range(num_keys):
+            if self.current_index >= num_keys:
+                self.current_index = 0
+            current_key = self.api_keys[self.current_index]
+            self.current_index += 1
+
+            if current_key not in api_key_blacklist:
+                return current_key
+
+        logger.error("所有API key都已耗尽或被黑名单，请重新配置或稍后重试")
+        return None 
+
+    def show_all_keys(self):
+        logger.info(f"当前可用API key个数: {len(self.api_keys)} ")
         for i, api_key in enumerate(self.api_keys):
             logger.info(f"API Key{i}: {api_key[:11]}...")
 
-    def get_available_key(self):
-        if self.current_index >= len(self.api_keys):
-            self.current_index = 0
-        current_key = self.api_keys[self.current_index]
-        self.current_index += 1
-        return current_key
+    def blacklist_key(self, key):
+        logger.warning(f"{key[:11]} → 加入黑名单 {api_key_blacklist_duration} 秒")
+        api_key_blacklist.add(key)
+
+        scheduler.add_job(lambda: api_key_blacklist.discard(key), 'date', run_date=datetime.now() + timedelta(seconds=api_key_blacklist_duration))
 
 key_manager = APIKeyManager()
+key_manager.show_all_keys()
 current_api_key = key_manager.get_available_key()
-logger.info(f"Current API key: {current_api_key[:11]}...")
+
+def switch_api_key():
+    global current_api_key
+    key = key_manager.get_available_key()
+    if key:
+      current_api_key = key
+      logger.info(f"API key 替换为 → {current_api_key[:11]}...")
+    else:
+      logger.error("API key 替换失败，所有API key都已耗尽或被黑名单，请重新配置或稍后重试")
+
+logger.info(f"当前 API key: {current_api_key[:11]}...")
 
 GEMINI_MODELS = [
-    {"id": "gemini-pro"},
-    {"id": "gemini-pro-vision"},
-    {"id": "gemini-1.0-pro"},
-    {"id": "gemini-1.0-pro-vision"},
     {"id": "gemini-1.5-pro-002"},
+    {"id": "gemini-1.5-pro-latest"},
+    {"id": "gemini-1.5-pro-exp-0827"},
+    {"id": "learnlm-1.5-pro-experimental"},
     {"id": "gemini-exp-1114"},
     {"id": "gemini-exp-1121"},
     {"id": "gemini-exp-1206"},
     {"id": "gemini-2.0-flash-exp"},
     {"id": "gemini-2.0-flash-thinking-exp-1219"},
-    {"id": "gemini-2.0-pro-exp"},
+    {"id": "gemini-2.0-pro-exp"}
 ]
+
 @app.route('/')
 def index():
-    main_content = "Moonfanz Gemini"
+    main_content = "Moonfanz Reminiproxy"
     html_template = """
 <!DOCTYPE html>
 <html>
@@ -109,146 +130,195 @@ function copyLink(event) {
 </script>
 </head>
 <body>
-{{ main_content }}<br/><br/>完全开源、免费且禁止商用<br/><br/>点击复制反向代理: <a href="v1" onclick="copyLink(event)">Copy Link</a><br/>聊天来源选择"自定义(兼容 OpenAI)"<br/>将复制的网址填入到自定义端点<br/>将环境设置中的password填入自定义API秘钥<br/><br/><br/>
+{{ main_content }}<br/><br/>完全开源、免费且禁止商用<br/><br/>点击复制反向代理: <a href="v1" onclick="copyLink(event)">Copy Link</a><br/>聊天来源选择"自定义(兼容 OpenAI)"<br/>将复制的网址填入到自定义端点<br/>将设置password填入自定义API秘钥<br/><br/><br/>
 </body>
 </html>
     """
     return render_template_string(html_template, main_content=main_content)
 
+def is_within_rate_limit(api_key):
+    now = datetime.now()
+    if api_key not in request_counts:
+        request_counts[api_key] = deque()
+
+    while request_counts[api_key] and request_counts[api_key][0] < now - timedelta(seconds=LIMIT_WINDOW):
+        request_counts[api_key].popleft()
+
+    return len(request_counts[api_key]) < MAX_REQUESTS
+
+def increment_request_count(api_key):
+    now = datetime.now()
+    if api_key not in request_counts:
+        request_counts[api_key] = deque()
+    request_counts[api_key].append(now)
+
+def handle_api_error(error, attempt, stream=False):
+    if attempt > MAX_RETRIES:
+        logger.error(f"{MAX_RETRIES} 次尝试后仍然失败，请修改预设或输入")
+        return False, jsonify({
+                'error': {
+                    'message': f"{MAX_RETRIES} 次尝试后仍然失败，请修改预设或输入",
+                    'type': 'max_retries_exceeded'
+                }
+        })
+
+    if isinstance(error, InvalidArgument):
+        logger.error(f"{current_api_key[:11]} → 无效，可能被删除或过期")
+        key_manager.blacklist_key(current_api_key)
+        switch_api_key()
+        return False, None
+
+    elif isinstance(error, (ResourceExhausted, Aborted, InternalServerError, ServiceUnavailable)):
+        delay = min(RETRY_DELAY * (2 ** attempt), MAX_RETRY_DELAY)
+        if isinstance(error, ResourceExhausted):
+          logger.warning(f"{current_api_key[:11]} → 超过限额；{delay} 秒后重试...")
+        else:
+          logger.warning(f"{current_api_key[:11]} → 未知错误↙ \n{type(error).__name__}\n{delay} 秒后重试...")
+        time.sleep(delay)
+        if isinstance(error, (ResourceExhausted)):
+            key_manager.blacklist_key(current_api_key)
+            switch_api_key()
+        return False, None
+
+    elif isinstance(error, generation_types.StopCandidateException):
+        logger.warning(f"输出截断")
+        switch_api_key()
+        return False, None
+
+    else:
+        logger.error(f"未知错误↙\n {error}")
+        return False, None
+
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 def chat_completions():
-    global current_api_key
     is_authenticated, auth_error, status_code = func.authenticate_request(request)
     if not is_authenticated:
         return auth_error if auth_error else jsonify({'error': '未授权'}), status_code if status_code else 401
-    try:
-        request_data = request.get_json()
-        messages = request_data.get('messages', [])
-        model = request_data.get('model', 'gemini-exp-1206')  
-        temperature = request_data.get('temperature', 1)
-        max_tokens = request_data.get('max_tokens', 8192)
-        stream = request_data.get('stream', False)
 
-        logger.info(f"\n{model} [r] -> {current_api_key[:11]}...")
+    request_data = request.get_json()
+    messages = request_data.get('messages', [])
+    model = request_data.get('model', 'gemini-2.0-flash-exp')
+    temperature = request_data.get('temperature', 1)
+    max_tokens = request_data.get('max_tokens', 8192)
+    stream = request_data.get('stream', False)
 
-        gemini_history, user_message, error_response = func.process_messages_for_gemini(messages)
+    logger.info(f"\n{model} [r] → {current_api_key[:11]}...")
 
-        if error_response:
-            print(error_response)
+    gemini_history, user_message, error_response = func.process_messages_for_gemini(messages)
 
-        genai.configure(api_key=current_api_key)
+    if error_response:
+      logger.error(f"Error in processing messages: {error_response}")
+      return jsonify(error_response), 400
 
-        generation_config = {
-            "temperature": temperature,
-            "max_output_tokens": max_tokens
-        }
+    def do_request(current_api_key, attempt):
+        if not is_within_rate_limit(current_api_key):
+            logger.warning(f"{current_api_key[:11]} → 暂时超过限额")
+            key_manager.blacklist_key(current_api_key)
+            switch_api_key()
+            return False, None
 
-        gen_model = genai.GenerativeModel(
-            model_name=model,
-            generation_config=generation_config,
-            safety_settings=safety_settings
-        )
+        increment_request_count(current_api_key)
 
-        if stream:
+        gen_model = func.get_gen_model(current_api_key, model, temperature, max_tokens)
 
+        try:
             if gemini_history:
                 chat_session = gen_model.start_chat(history=gemini_history)
-                response = chat_session.send_message(user_message, stream=True)
+                response = chat_session.send_message(user_message, stream=stream)
             else:
-                response = gen_model.generate_content(user_message, stream=True)
-
-            def generate():
-                try:
-                    for chunk in response:
-                        if chunk.text:
-                            data = {
-                                'choices': [
-                                    {
-                                        'delta': {
-                                            'content': chunk.text
-                                        },
-                                        'finish_reason': None,
-                                        'index': 0
-                                    }
-                                ],
-                                'object': 'chat.completion.chunk'
-                            }
-
-                            yield f"data: {json.dumps(data)}\n\n"
+                response = gen_model.generate_content(user_message, stream=stream)
+            return True, response
+        except Exception as e:
+            return handle_api_error(e, attempt, stream)
+
+    def generate(response):
+        try:
+            for chunk in response:
+                if chunk.text:
                     data = {
                         'choices': [
                             {
-                                'delta': {},
-                                'finish_reason': 'stop',
+                                'delta': {
+                                    'content': chunk.text
+                                },
+                                'finish_reason': None,
                                 'index': 0
                             }
                         ],
                         'object': 'chat.completion.chunk'
                     }
-                    logger.info(f"200!")
-                    yield f"data: {json.dumps(data)}\n\n"
-                except Exception as e:
-                    
-                    logger.error(f"Error during streaming: {str(e)}")
-                    current_api_key = key_manager.get_available_key()
-                    logger.info(f"API KEY Switched -> {current_api_key[:11]}...")
-                    data = {
-                        'error': {
-                            'message': str(e),
-                            'type': 'internal_server_error'
-                        }
-                    }
                     yield f"data: {json.dumps(data)}\n\n"
 
-            return Response(stream_with_context(generate()), mimetype='text/event-stream')
-        else:
+            
+            yield "data: [DONE]\n\n"
+            logger.info(f"200!")
 
-            if gemini_history:
-                chat_session = gen_model.start_chat(history=gemini_history)
-                response = chat_session.send_message(user_message)
-            else:
-                response = gen_model.generate_content(user_message)
-
-            try:
-                text_content = response.candidates[0].content.parts[0].text
-
-            except (AttributeError, IndexError, TypeError) as e:
-                logger.error(f"Error getting text content: {str(e)}")
-
-                text_content = "Error: Unable to get text content." 
-
-            response_data = {
-                'id': 'chatcmpl-xxxxxxxxxxxx',
-                'object': 'chat.completion',
-                'created': int(datetime.now().timestamp()),
-                'model': model,
-                'choices': [{
-                    'index': 0,
-                    'message': {
-                        'role': 'assistant',
-                        'content': text_content
-                    },
-                    'finish_reason': 'stop'
-                }],
-                'usage':{ 
-                    'prompt_tokens': 0,
-                    'completion_tokens': 0,
-                    'total_tokens': 0
+        except Exception as e:
+            switch_api_key()
+            logger.error(f"流式输出时截断")
+            error_data = {
+                'error': {
+                    'message': '流式输出时截断，请关闭流式输出或修改你的输入',
+                    'type': 'internal_server_error'
                 }
             }
-            logger.info(f"200!")
-            return jsonify(response_data)
+            yield f"data: {json.dumps(error_data)}\n\n"
+            yield "data: [DONE]\n\n"
+
+    attempt = 0
+    success = False
+    response = None
+
+    while attempt < MAX_RETRIES and not success:
+        attempt += 1
+        logger.info(f"尝试中... 第 {attempt}/{MAX_RETRIES} 次尝试")
+        success, response = do_request(current_api_key, attempt)
 
-    except Exception as e:
-        logger.error(f"Error in chat completions: {str(e)}")
-        current_api_key = key_manager.get_available_key()
-        logger.info(f"API KEY Switched -> {current_api_key[:11]}...")
-        return jsonify({
+    if not success:
+        logger.error(f" {MAX_RETRIES} 次尝试均失败，请修改预设")
+        response = {
             'error': {
-                'message': str(e),
-                'type': 'invalid_request_error'
+                'message': f' {MAX_RETRIES} 次尝试均失败，请修改预设',
+                'type': 'internal_server_error'
             }
-        }), 500
+        }
+        return jsonify(response), 500 if response is not None else 503
+
+    if stream:
+        return Response(stream_with_context(generate(response)), mimetype='text/event-stream')
+    else:
+        try:
+            text_content = response.text
+        except (AttributeError, IndexError, TypeError) as e:
+              logger.error(f"处理响应时出错↙\n{e}")
+              return jsonify({
+                  'error': {
+                      'message': '处理响应时出错',
+                      'type': 'response_processing_error'
+                  }
+              }), 500
+
+        response_data = {
+            'id': 'chatcmpl-xxxxxxxxxxxx',  
+            'object': 'chat.completion',
+            'created': int(datetime.now().timestamp()),
+            'model': model,
+            'choices': [{
+                'index': 0,
+                'message': {
+                    'role': 'assistant',
+                    'content': text_content
+                },
+                'finish_reason': 'stop'
+            }],
+            'usage': {
+                'prompt_tokens': 0,  
+                'completion_tokens': 0,
+                'total_tokens': 0
+            }
+        }
+        logger.info(f"200!")
+        return jsonify(response_data)
 
 @app.route('/hf/v1/models', methods=['GET'])
 def list_models():
@@ -265,7 +335,12 @@ def keep_alive():
 
 if __name__ == '__main__':
     scheduler = BackgroundScheduler()
-    scheduler.add_job(keep_alive, 'interval', hours = 12)
     
+    scheduler.add_job(keep_alive, 'interval', hours=12)
     scheduler.start()
+
+    logger.info(f"最大尝试次数/Max retries: {MAX_RETRIES}")
+    logger.info(f"最大请求次数/Max requests: {MAX_REQUESTS}")
+    logger.info(f"请求限额窗口/Limit window: {LIMIT_WINDOW} seconds")
+
     app.run(debug=True, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))

func.py

@@ -6,8 +6,33 @@ import logging
 import json
 import re
 import os
+import requests
+import google.generativeai as genai
 logger = logging.getLogger(__name__)
 
+
+request_counts = {}
+
+# 核心优势
+safety_settings = [
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "BLOCK_NONE"
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "BLOCK_NONE"
+    },
+]
+
 password = os.environ['password']
 
 def authenticate_request(request):
@@ -15,7 +40,7 @@ def authenticate_request(request):
 
     if not auth_header:
         return False, jsonify({'error': '缺少Authorization请求头'}), 401
-    
+
     try:
         auth_type, pass_word = auth_header.split(' ', 1)
     except ValueError:
@@ -29,8 +54,23 @@ def authenticate_request(request):
 
     return True, None, None
 
+def get_gen_model(api_key, model, temperature, max_tokens):
+    genai.configure(api_key=api_key)
+
+    generation_config = {
+        "temperature": temperature,
+        "max_output_tokens": max_tokens
+        }
+
+    gen_model = genai.GenerativeModel(
+        model_name=model,
+        generation_config=generation_config,
+        safety_settings=safety_settings
+    )
+    return gen_model
+
 def process_messages_for_gemini(messages):
-    
+
     gemini_history = []
     for message in messages:
         role = message.get('role')

requirements.txt

@@ -5,5 +5,4 @@ Werkzeug==2.0.3
 google==3.0.0
 google-generativeai==0.8.3
 pillow==10.4.0
-apscheduler
-Flask-HTTPAuth==4.8.0
+apscheduler