Upload 2 files

app.py

@@ -4,12 +4,10 @@ import google.generativeai as genai
 import json
 from datetime import datetime
 import os
-import base64
 from termcolor import colored
 import logging
-from PIL import Image
-from io import BytesIO
-import re
+import func
+
 os.environ['TZ'] = 'Asia/Shanghai'
 app = Flask(__name__)
 if 'API_KEYS' not in os.environ:
@@ -63,78 +61,31 @@ current_api_key = key_manager.get_available_key()
 logger.info(f"Current API key: {current_api_key}")
 # 模型列表
 GEMINI_MODELS = [
-    {"id": "gemini-pro", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-pro-vision", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-1.0-pro", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-1.0-pro-vision", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-1.5-pro-002", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-exp-1114", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-exp-1121", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-exp-1206", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-2.0-flash-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-2.0-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
-    {"id": "gemini-2.0-pro-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"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-exp-1114"},
+    {"id": "gemini-exp-1121"},
+    {"id": "gemini-exp-1206"},
+    {"id": "gemini-2.0-flash-exp"},
+    {"id": "gemini-2.0-exp"},
+    {"id": "gemini-2.0-pro-exp"},
 ]
-def authenticate_request(request):
-    auth_header = request.headers.get('Authorization')
-    hf_api_key = os.environ.get('HF_API_KEY').split(',')
-
-    if not auth_header:
-        return False, jsonify({'error': 'Authorization header is missing'}), 401
-    
-    try:
-        auth_type, api_key = auth_header.split(' ', 1)
-    except ValueError:
-        return False, jsonify({'error': 'Invalid Authorization header format'}), 401
-
-    if auth_type.lower() != 'bearer':
-        return False, jsonify({'error': 'Authorization type must be Bearer'}), 401
-
-    if api_key not in hf_api_key:
-        return False, jsonify({'error': 'Unauthorized'}), 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
-    )
-
-    return json.loads(sanitized_request_data_str)
 
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 def chat_completions():
     global current_api_key
-    is_authenticated, auth_error, status_code = authenticate_request(request)
+    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 = sanitize_request_data(request_data)
+        r_data = func.sanitize_request_data(request_data)
         logger.info(json.dumps(r_data, indent=4, ensure_ascii=False), extra={"color": "green"})
         messages = request_data.get('messages', [])
-        model = request_data.get('model', 'gemini-exp-1206')  # 默认改为 gemini-pro-vision 用于处理图文
+        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)
@@ -142,67 +93,12 @@ def chat_completions():
         logger.info(colored(f"\n{model} [r] -> {current_api_key[:11]}...", 'yellow'))
 
         # 将 OpenAI 格式的消息转换为 Gemini 格式
-        gemini_history = []
-        for message in messages:
-            role = message.get('role')
-            content = message.get('content')
-
-            if isinstance(content, str): # 纯文本
-                if role == 'system':
-                    gemini_history.append({"role": "user", "parts": [content]})
-                elif role == 'user':
-                    gemini_history.append({"role": "user", "parts": [content]})
-                elif role == 'assistant':
-                    gemini_history.append({"role": "model", "parts": [content]})
-            elif isinstance(content, list):  # 图文
-                parts = []
-                for item in content:
-                    if item.get('type') == 'text':
-                        parts.append(item.get('text'))
-                    elif item.get('type') == 'image_url':
-                        image_data = item.get('image_url', {}).get('url', '')
-                        if image_data.startswith('data:image/'):  # 修改判断条件
-                            try:
-                                # 提取 base64 编码和图片类型
-                                image_type = image_data.split(';')[0].split('/')[1].upper()  # 提取图片类型并转为大写
-                                base64_image = image_data.split(';base64,')[1]
-
-                                image = Image.open(BytesIO(base64.b64decode(base64_image)))
-
-                                # 将图片转换为 RGB 模式
-                                if image.mode != 'RGB':
-                                    image = image.convert('RGB')
-
-                                # 压缩图像
-                                if image.width > 2048 or image.height > 2048:
-                                    image.thumbnail((2048, 2048))
-
-                                output_buffer = BytesIO()
-                                image.save(output_buffer, format=image_type)  # 使用原始图片类型保存
-                                output_buffer.seek(0)
-                                parts.append(image)  # 直接添加 image 对象
-                            except Exception as e:
-                                logger.error(f"Error processing image: {e}")
-                                return jsonify({'error': 'Invalid image data'}), 400
-                        else:
-                            return jsonify({'error': 'Invalid image URL format'}), 400
-
-
-                # 根据 role 添加到 gemini_history
-                if role in ['user', 'system']:
-                    gemini_history.append({"role": "user", "parts": parts}) # 保持 parts 的原有顺序
-                elif role == 'assistant':
-                    gemini_history.append({"role": "model", "parts": parts})
-                else:
-                    return jsonify({'error': f'Invalid role: {role}'}), 400
-
-        # 用户最后一条消息
-        if gemini_history:
-            user_message = gemini_history[-1]
-            gemini_history = gemini_history[:-1]  # 历史记录不包含最后一条消息
-        else:
-            user_message = {"role": "user", "parts": [""]}
-
+        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 = {
@@ -305,10 +201,10 @@ def chat_completions():
             }
             logger.info(colored(f"Generation Success", 'green'))
             return jsonify(response_data)
-            
+
     except Exception as e:
         logger.error(f"Error in chat completions: {str(e)}")
-        
+
         return jsonify({
             'error': {
                 'message': str(e),
@@ -319,10 +215,11 @@ def chat_completions():
         current_api_key = key_manager.get_available_key()
         logger.info(colored(f"API KEY Switched -> {current_api_key[:11]}...", 'aqua'))
 
-
-
 @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)
 

func.py

@@ -0,0 +1,152 @@
+from io import BytesIO
+import base64
+from PIL import Image
+from flask import jsonify
+import logging
+import json
+import re
+import os
+GEMINI_MODELS = [
+    {"id": "gemini-pro", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-pro-vision", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-1.0-pro", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-1.0-pro-vision", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-1.5-pro-002", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-exp-1114", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-exp-1121", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-exp-1206", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-2.0-flash-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-2.0-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
+    {"id": "gemini-2.0-pro-exp", "object": "model", "created": 1700000000, "owned_by": "google"},
+]
+
+logger = logging.getLogger(__name__)
+
+def authenticate_request(request):
+    auth_header = request.headers.get('Authorization')
+    hf_api_key = os.environ.get('HF_API_KEY').split(',')
+
+    if not auth_header:
+        return False, jsonify({'error': 'Authorization header is missing'}), 401
+    
+    try:
+        auth_type, api_key = auth_header.split(' ', 1)
+    except ValueError:
+        return False, jsonify({'error': 'Invalid Authorization header format'}), 401
+
+    if auth_type.lower() != 'bearer':
+        return False, jsonify({'error': 'Authorization type must be Bearer'}), 401
+
+    if api_key not in hf_api_key:
+        return False, jsonify({'error': 'Unauthorized'}), 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
+    )
+
+    return json.loads(sanitized_request_data_str)
+
+def process_messages_for_gemini(messages):
+    """
+    Processes a list of messages to construct the history format expected by Gemini,
+    handling both text and image content.
+
+    Args:
+        messages: A list of message dictionaries. Each dictionary should have a 'role' key
+                  ('system', 'user', or 'assistant') and a 'content' key.
+                  The 'content' can be either a string (for text messages) or a list
+                  (for multi-modal messages).
+                  For multi-modal messages, the list should contain dictionaries with 'type'
+                  ('text' or 'image_url') and corresponding 'text' or 'image_url' keys.
+
+    Returns:
+        A tuple containing:
+            - gemini_history: A list of dictionaries formatted for the Gemini API, excluding the last user message.
+            - user_message: The last user message from the input, or a default empty user message if the input is empty.
+            - error_response: A Flask error response (jsonify, status_code) if an error occurred, otherwise None.
+    """
+    gemini_history = []
+    for message in messages:
+        role = message.get('role')
+        content = message.get('content')
+
+        if isinstance(content, str):  # 纯文本
+            if role == 'system':
+                gemini_history.append({"role": "user", "parts": [content]})
+            elif role == 'user':
+                gemini_history.append({"role": "user", "parts": [content]})
+            elif role == 'assistant':
+                gemini_history.append({"role": "model", "parts": [content]})
+        elif isinstance(content, list):  # 图文
+            parts = []
+            for item in content:
+                if item.get('type') == 'text':
+                    parts.append(item.get('text'))
+                elif item.get('type') == 'image_url':
+                    image_data = item.get('image_url', {}).get('url', '')
+                    if image_data.startswith('data:image/'):
+                        try:
+                            # 提取 base64 编码和图片类型
+                            image_type = image_data.split(';')[0].split('/')[1].upper()  # 提取图片类型并转为大写
+                            base64_image = image_data.split(';base64,')[1]
+
+                            image = Image.open(BytesIO(base64.b64decode(base64_image)))
+
+                            # 将图片转换为 RGB 模式
+                            if image.mode != 'RGB':
+                                image = image.convert('RGB')
+
+                            # 压缩图像
+                            if image.width > 2048 or image.height > 2048:
+                                image.thumbnail((2048, 2048))
+
+                            output_buffer = BytesIO()
+                            image.save(output_buffer, format=image_type)  # 使用原始图片类型保存
+                            output_buffer.seek(0)
+                            parts.append(image)
+                        except Exception as e:
+                            logger.error(f"Error processing image: {e}")
+                            return [], None, (jsonify({'error': 'Invalid image data'}), 400)
+                    else:
+                        return [], None, (jsonify({'error': 'Invalid image URL format'}), 400)
+
+            # 根据 role 添加到 gemini_history
+            if role in ['user', 'system']:
+                gemini_history.append({"role": "user", "parts": parts})
+            elif role == 'assistant':
+                gemini_history.append({"role": "model", "parts": parts})
+            else:
+                return [], None, (jsonify({'error': f'Invalid role: {role}'}), 400)
+
+    # 用户最后一条消息
+    if gemini_history:
+        user_message = gemini_history[-1]
+        gemini_history = gemini_history[:-1]  # 历史记录不包含最后一条消息
+    else:
+        user_message = {"role": "user", "parts": [""]}
+
+    return gemini_history, user_message, None