Upload 4 files

app.py

@@ -1,272 +1,334 @@
-
-from flask import Flask, request, jsonify, Response, stream_with_context, render_template, session, redirect, url_for
-from apscheduler.schedulers.background import BackgroundScheduler
+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
-from termcolor import colored
 import logging
 import func
-import requests
+from datetime import datetime, timedelta
+from apscheduler.schedulers.background import BackgroundScheduler
 import time
+import requests
+from collections import deque
+
 
 os.environ['TZ'] = 'Asia/Shanghai'
+
 app = Flask(__name__)
-if not os.environ.get('TERM'):
-    os.environ['TERM'] = 'xterm'
+
 app.secret_key = os.urandom(24)
 
-ProxyPasswords = os.environ.get('ProxyPasswords').split(',')
 
 formatter = logging.Formatter('%(message)s') 
-ADMIN_PASSWORD = os.environ.get('AdminPassword')
-
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO) 
-
+logger.setLevel(logging.INFO)
 handler = logging.StreamHandler()
 handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+MAX_RETRIES = int(os.environ.get('MaxRetries', 3))
+MAX_REQUESTS = int(os.environ.get('MaxRequests', 2))
+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
 
-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"
-    },
-]
 class APIKeyManager:
     def __init__(self):
-        self.api_keys = os.environ.get('API_KEYS').split(',')
-        self.current_index = 0  
-        
-            
+        self.api_keys = os.environ.get('KeyArray').split(',')
+        self.current_index = 0
+
     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
+        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 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]}...")
-genai.configure(api_key=current_api_key)
-models = genai.list_models()
-logger.info("Available models:")
-for model in models:
-    logger.info(f"- {model.name}: {model.supported_generation_methods}")
+
+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("/", methods=["GET", "POST"])
-def login():
-    if request.method == "POST":
-        password = request.form.get("password")
-        if password == ADMIN_PASSWORD:
-            session["logged_in"] = True
-            return redirect(url_for("manage_keys"))
+@app.route('/')
+def index():
+    main_content = "Moonfanz Reminiproxy"
+    html_template = """
+<!DOCTYPE html>
+<html>
+<head>
+<meta charset="utf-8">
+<script>
+function copyToClipboard(text) {
+  var textarea = document.createElement("textarea");
+  textarea.textContent = text;
+  textarea.style.position = "fixed";
+  document.body.appendChild(textarea);
+  textarea.select();
+  try {
+    return document.execCommand("copy");
+  } catch (ex) {
+    console.warn("Copy to clipboard failed.", ex);
+    return false;
+  } finally {
+    document.body.removeChild(textarea);
+  }
+}
+function copyLink(event) {
+  event.preventDefault();
+  const url = new URL(window.location.href);
+  const link = url.protocol + '//' + url.host + '/hf/v1';
+  copyToClipboard(link);
+  alert('链接已复制: ' + link);
+}
+</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/>
+</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()
+
+    if len(request_counts[api_key]) >= MAX_REQUESTS:
+        earliest_request_time = request_counts[api_key][0]
+        wait_time = (earliest_request_time + timedelta(seconds=LIMIT_WINDOW)) - now
+        return False, wait_time.total_seconds()
+    else:
+        return True, 0
+
+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]} → 429 官方资源耗尽 → {delay} 秒后重试...")
         else:
-            return render_template("login.html", error="Incorrect password")
-    return render_template("login.html", error=None)
-
-@app.route("/manage", methods=["GET", "POST"])
-def manage_keys():
-    if not session.get("logged_in"):
-        return redirect(url_for("login"))
-
-    if request.method == "POST":
-        action = request.form.get("action")
-        if action == "add":
-            new_key = request.form.get("new_key")
-            if new_key:
-                ProxyPasswords.append(new_key)
-        elif action == "delete":
-            key_to_delete = request.form.get("key_to_delete")
-            if key_to_delete in ProxyPasswords:
-                ProxyPasswords.remove(key_to_delete)
-
-    return render_template("manage.html", keys=ProxyPasswords)
-
-@app.route("/logout")
-def logout():
-    session.pop("logged_in", None)
-    return redirect(url_for("login"))
+            logger.warning(f"{current_api_key[:11]} → 未知错误↙ {delay} 秒后重试...\n{type(error).__name__}\n")
+        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"AI输出内容被Gemini官方阻挡，代理没有得到有效回复")
+        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': 'Unauthorized'}), status_code if status_code else 401
-    try:
-        request_data = request.get_json()
-        r_data = func.sanitize_request_data(request_data)
-        r_data = json.dumps(r_data, indent=4, ensure_ascii=False).replace('\\n', '\n')
-        os.system('cls' if os.name == 'nt' else 'clear')
-        logger.info(r_data)
-        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(colored(f"\n{model} [r] -> {current_api_key[:11]}...", 'yellow'))
-    
-        # 将 OpenAI 格式的消息转换为 Gemini 格式
-        gemini_history, user_message, error_response = func.process_messages_for_gemini(messages)
-        
-        if error_response:
-            # 处理错误
-            print(error_response)
-            
-        genai.configure(api_key=current_api_key)
-
-        generation_config = {
-            "temperature": temperature,
-            "max_output_tokens": max_tokens
-        }
+        return auth_error if auth_error else jsonify({'error': '未授权'}), status_code if status_code else 401
+
+    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)
+
+    logger.info(f"\n{model} [r] → {current_api_key[:11]}...")
+
+    gemini_history, user_message, error_response = func.process_messages_for_gemini(messages)
+
+    if error_response:
+      logger.error(f"处理输入消息时出错↙\n {error_response}")
+      return jsonify(error_response), 400
 
-        gen_model = genai.GenerativeModel(
-            model_name=model,
-            generation_config=generation_config,
-            safety_settings=safety_settings
-        )
+    def do_request(current_api_key, attempt):
+        isok, time = is_within_rate_limit(current_api_key)
+        if not isok:
+            logger.warning(f"{current_api_key[:11]} → 暂时超过限额，该API key将在 {time} 秒后启用...")
+            switch_api_key()
+            return False, None
 
-        
-        if stream:
-            # 流式响应
+        increment_request_count(current_api_key)
+
+        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'
                     }
-                    
-                    yield f"data: {json.dumps(data)}\n\n"
-                except Exception as e:
-                    logger.error(f"Error during streaming: {str(e)}")
-                    
-                    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:
-            # 调用 API
-            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(colored(f"Error getting text content: {str(e)}",'red'))
-
-                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
+            yield "data: [DONE]\n\n"
+            logger.info(f"200!")
+
+        except Exception as e:
+            logger.error(f"流式输出时截断，请关闭流式输出或修改你的输入")
+            error_data = {
+                'error': {
+                    'message': '流式输出时截断，请关闭流式输出或修改你的输入',
+                    'type': 'internal_server_error'
                 }
             }
-            logger.info(colored(f"Generation Success", 'green'))
-            return jsonify(response_data)
+            yield f"data: {json.dumps(error_data)}\n\n"
+            yield "data: [DONE]\n\n"
 
-    except Exception as e:
-        logger.error(f"Error in chat completions: {str(e)}")
-        return jsonify({
+    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)
+
+    if not success:
+        logger.error(f" {MAX_RETRIES} 次尝试均失败，请调整配置或向Moonfanz反馈")
+        response = {
             'error': {
-                'message': str(e),
-                'type': 'invalid_request_error'
+                'message': f' {MAX_RETRIES} 次尝试均失败，请调整配置或向Moonfanz反馈',
+                'type': 'internal_server_error'
+            }
+        }
+        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"处理AI返回消息时出错↙\n{e}")
+              return jsonify({
+                  'error': {
+                      'message': '处理AI返回消息时出错',
+                      '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
             }
-        }), 500
-    finally:
-        current_api_key = key_manager.get_available_key()
-        logger.info(colored(f"API KEY Switched -> {current_api_key[:11]}...", 'aqua'))
+        }
+        logger.info(f"200!")
+        return jsonify(response_data)
 
 @app.route('/hf/v1/models', methods=['GET'])
 def list_models():
-    is_authenticated, auth_error, status_code = func.authenticate_request(request)
-    if not is_authenticated:
-        return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
     response = {"object": "list", "data": GEMINI_MODELS}
     return jsonify(response)
 
 def keep_alive():
-    """ 定期向应用自身发送请求，保持活跃 """
     try:
         response = requests.get("http://127.0.0.1:7860/", timeout=10)
         response.raise_for_status()  
@@ -276,8 +338,12 @@ def keep_alive():
 
 if __name__ == '__main__':
     scheduler = BackgroundScheduler()
-    # 设置定时任务，每 12 小时执行一次 keep_alive 函数
-    scheduler.add_job(keep_alive, 'interval', hours = 12)
-    
+
+    scheduler.add_job(keep_alive, 'interval', hours=12)
     scheduler.start()
-    app.run(debug=True, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))
+
+    logger.info(f"最大尝试次数/MaxRetries: {MAX_RETRIES}")
+    logger.info(f"最大请求次数/MaxRequests: {MAX_REQUESTS}")
+    logger.info(f"请求限额窗口/LimitWindow: {LIMIT_WINDOW} 秒")
+
+    app.run(debug=True, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))

func.py

@@ -1,92 +1,75 @@
 from io import BytesIO
-import os
+import base64
+from PIL import Image
 from flask import jsonify
 import logging
 import json
 import re
+import os
+import requests
+import google.generativeai as genai
 logger = logging.getLogger(__name__)
 
-ProxyPasswords = os.environ.get('ProxyPasswords').split(',')
+
+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):
     auth_header = request.headers.get('Authorization')
 
     if not auth_header:
-        return False, jsonify({'error': 'Authorization header is missing'}), 401
-    
+        return False, jsonify({'error': '缺少Authorization请求头'}), 401
+
     try:
-        auth_type, api_key = auth_header.split(' ', 1)
+        auth_type, pass_word = auth_header.split(' ', 1)
     except ValueError:
-        return False, jsonify({'error': 'Invalid Authorization header format'}), 401
+        return False, jsonify({'error': 'Authorization请求头格式错误'}), 401
 
     if auth_type.lower() != 'bearer':
-        return False, jsonify({'error': 'Authorization type must be Bearer'}), 401
+        return False, jsonify({'error': 'Authorization类型必须为Bearer'}), 401
 
-    if api_key not in ProxyPasswords:
-        return False, jsonify({'error': 'Unauthorized'}), 401
+    if pass_word != password:
+        return False, jsonify({'error': '未授权'}), 401
 
     return True, None, None
 
-def sanitize_request_data(request_data):
-    """
-    从请求数据中删除base64编码的数据。
-
-    Args:
-        request_data: 包含可能存在base64数据的字典。
-
-    Returns:
-        清理后的字典，其中base64数据被替换为"[Base64 Data Omitted]"。
-    """
-    
-    
-    def replace_base64(match):
-        # 替换base64数据为提示信息
-        return '"[Base64 Data Omitted]"'
-
-    
-    request_data_str = json.dumps(request_data)
-
-    # 使用正则表达式匹配base64数据，并替换为提示信息
-    sanitized_request_data_str = re.sub(
-        r'"(data:[^;]+;base64,)[^"]+"',
-        replace_base64,
-        request_data_str
-    )
+def get_gen_model(api_key, model, temperature, max_tokens):
+    genai.configure(api_key=api_key)
 
-    return json.loads(sanitized_request_data_str)
+    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 API所需的格式
-    这个函数处理消息列表并将其转换为Gemini API兼容的格式。它支持文本、图片和文件内容的处理。
-    参数:
-        messages (list): 包含对话消息的列表。每条消息应该是一个字典，包含'role'和'content'字段。
-            - role: 可以是 'system', 'user' 或 'assistant'
-            - content: 可以是字符串或包含多个内容项的列表
-    返回:
-        tuple: 包含三个元素:
-            - gemini_history (list): 转换后的历史消息列表
-            - user_message (dict): 最新的用户消息
-            - error (tuple or None): 如果有错误，返回错误响应；否则返回None
-    错误处理:
-        - 检查角色是否有效
-        - 验证图片URL格式
-        - 验证文件URL格式
-    示例消息格式:
-        文本消息:
-        {
-            'role': 'user',
-            'content': '你好'
-        }
-        多模态消息:
-        {
-            'role': 'user',
-            'content': [
-                {'type': 'text', 'text': '这是什么图片?'},
-                {'type': 'image_url', 'image_url': {'url': 'data:image/jpeg;base64,...'}}
-            ]
-        }
-    """
     gemini_history = []
     errors = []
     for message in messages:

requirements.txt

@@ -4,6 +4,5 @@ requests==2.26.0
 Werkzeug==2.0.3
 google==3.0.0
 google-generativeai==0.8.3
-termcolor==2.5.0
 pillow==10.4.0
-apscheduler
+apscheduler