Upload 10 files

app/gemini.py

@@ -0,0 +1,256 @@
+import requests
+import json
+import os
+import asyncio
+from app.models import ChatCompletionRequest, Message  # 相对导入
+from dataclasses import dataclass
+from typing import Optional, Dict, Any, List
+import httpx
+
+
+@dataclass
+class GeneratedText:
+    text: str
+    finish_reason: Optional[str] = None
+
+
+class ResponseWrapper:
+    def __init__(self, data: Dict[Any, Any]):  # 正确的初始化方法名
+        self._data = data
+        self._text = self._extract_text()
+        self._finish_reason = self._extract_finish_reason()
+        self._prompt_token_count = self._extract_prompt_token_count()
+        self._candidates_token_count = self._extract_candidates_token_count()
+        self._total_token_count = self._extract_total_token_count()
+        self._thoughts = self._extract_thoughts()
+        self._json_dumps = json.dumps(self._data, indent=4, ensure_ascii=False)
+
+    def _extract_thoughts(self) -> Optional[str]:
+        try:
+            for part in self._data['candidates'][0]['content']['parts']:
+                if 'thought' in part:
+                    return part['text']
+            return ""
+        except (KeyError, IndexError):
+            return ""
+
+    def _extract_text(self) -> str:
+        try:
+            for part in self._data['candidates'][0]['content']['parts']:
+                if 'thought' not in part:
+                    return part['text']
+            return ""
+        except (KeyError, IndexError):
+            return ""
+
+    def _extract_finish_reason(self) -> Optional[str]:
+        try:
+            return self._data['candidates'][0].get('finishReason')
+        except (KeyError, IndexError):
+            return None
+
+    def _extract_prompt_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('promptTokenCount')
+        except (KeyError):
+            return None
+
+    def _extract_candidates_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('candidatesTokenCount')
+        except (KeyError):
+            return None
+
+    def _extract_total_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('totalTokenCount')
+        except (KeyError):
+            return None
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    @property
+    def finish_reason(self) -> Optional[str]:
+        return self._finish_reason
+
+    @property
+    def prompt_token_count(self) -> Optional[int]:
+        return self._prompt_token_count
+
+    @property
+    def candidates_token_count(self) -> Optional[int]:
+        return self._candidates_token_count
+
+    @property
+    def total_token_count(self) -> Optional[int]:
+        return self._total_token_count
+
+    @property
+    def thoughts(self) -> Optional[str]:
+        return self._thoughts
+
+    @property
+    def json_dumps(self) -> str:
+        return self._json_dumps
+
+
+class GeminiClient:
+
+    AVAILABLE_MODELS = []
+    EXTRA_MODELS = os.environ.get("EXTRA_MODELS", "").split(",")
+
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+
+    async def stream_chat(self, request: ChatCompletionRequest, contents, safety_settings, system_instruction):
+        api_version = "v1alpha" if "think" in request.model else "v1beta"
+        url = f"https://generativelanguage.googleapis.com/{api_version}/models/{request.model}:streamGenerateContent?key={self.api_key}&alt=sse"
+        headers = {
+            "Content-Type": "application/json",
+        }
+        data = {
+            "contents": contents,
+            "generationConfig": {
+                "temperature": request.temperature,
+                "maxOutputTokens": request.max_tokens,
+            },
+            "safetySettings": safety_settings,
+        }
+        if system_instruction:
+            data["system_instruction"] = system_instruction
+
+        async with httpx.AsyncClient() as client:
+            async with client.stream("POST", url, headers=headers, json=data, timeout=600) as response:
+                buffer = b""  # 初始化 JSON 缓冲
+                async for line in response.aiter_lines():
+                    if line.startswith("data: "): 
+                        line = line[len("data: "):] 
+                    buffer += line.encode('utf-8')
+                    try:
+                        data = json.loads(buffer.decode('utf-8'))
+                        buffer = b""
+                        if 'candidates' in data and data['candidates']:
+                            candidate = data['candidates'][0]
+                            if 'content' in candidate:
+                                content = candidate['content']
+                                if 'parts' in content and content['parts']:
+                                    parts = content['parts']
+                                    text = ""
+                                    for part in parts:
+                                        if 'text' in part:
+                                            text += part['text']
+                                    finish_reason = candidate.get('finishReason')
+                                    if text:
+                                        yield text
+                    except json.JSONDecodeError:
+                        continue
+                    except Exception as e:
+                        print(f"Error parsing JSON: {e}")
+                        continue
+
+    def complete_chat(self, request: ChatCompletionRequest, contents, safety_settings, system_instruction):
+        api_version = "v1alpha" if "think" in request.model else "v1beta"
+        url = f"https://generativelanguage.googleapis.com/{api_version}/models/{request.model}:generateContent?key={self.api_key}"
+        headers = {
+            "Content-Type": "application/json",
+        }
+        data = {
+            "contents": contents,
+            "generationConfig": {
+                "temperature": request.temperature,
+                "maxOutputTokens": request.max_tokens,
+            },
+            "safetySettings": safety_settings,
+        }
+        if system_instruction:
+            data["system_instruction"] = system_instruction
+        response = requests.post(url, headers=headers, json=data)
+        response.raise_for_status()
+        return ResponseWrapper(response.json())
+
+    def convert_messages(self, messages, use_system_prompt=False):
+        gemini_history = []
+        errors = []
+        system_instruction_text = ""
+        is_system_phase = use_system_prompt
+        for i, message in enumerate(messages):
+            role = message.role
+            content = message.content
+
+            if isinstance(content, str):
+                if is_system_phase and role == 'system':
+                    if system_instruction_text:
+                        system_instruction_text += "\n" + content
+                    else:
+                        system_instruction_text = content
+                else:
+                    is_system_phase = False
+
+                    if role in ['user', 'system']:
+                        role_to_use = 'user'
+                    elif role == 'assistant':
+                        role_to_use = 'model'
+                    else:
+                        errors.append(f"Invalid role: {role}")
+                        continue
+
+                    if gemini_history and gemini_history[-1]['role'] == role_to_use:
+                        gemini_history[-1]['parts'].append({"text": content})
+                    else:
+                        gemini_history.append(
+                            {"role": role_to_use, "parts": [{"text": content}]})
+            elif isinstance(content, list):
+                parts = []
+                for item in content:
+                    if item.get('type') == 'text':
+                        parts.append({"text": item.get('text')})
+                    elif item.get('type') == 'image_url':
+                        image_data = item.get('image_url', {}).get('url', '')
+                        if image_data.startswith('data:image/'):
+                            try:
+                                mime_type, base64_data = image_data.split(';')[
+                                    0].split(':')[1], image_data.split(',')[1]
+                                parts.append({
+                                    "inline_data": {
+                                        "mime_type": mime_type,
+                                        "data": base64_data
+                                    }
+                                })
+                            except (IndexError, ValueError):
+                                errors.append(
+                                    f"Invalid data URI for image: {image_data}")
+                        else:
+                            errors.append(
+                                f"Invalid image URL format for item: {item}")
+
+                if parts:
+                    if role in ['user', 'system']:
+                        role_to_use = 'user'
+                    elif role == 'assistant':
+                        role_to_use = 'model'
+                    else:
+                        errors.append(f"Invalid role: {role}")
+                        continue
+                    if gemini_history and gemini_history[-1]['role'] == role_to_use:
+                        gemini_history[-1]['parts'].extend(parts)
+                    else:
+                        gemini_history.append(
+                            {"role": role_to_use, "parts": parts})
+        if errors:
+            return errors
+        else:
+            return gemini_history, {"parts": [{"text": system_instruction_text}]}
+
+    @staticmethod
+    async def list_available_models(api_key) -> list:
+        url = "https://generativelanguage.googleapis.com/v1beta/models?key={}".format(
+            api_key)
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url)
+            response.raise_for_status()
+            data = response.json()
+            models = [model["name"] for model in data.get("models", [])]
+            models.extend(GeminiClient.EXTRA_MODELS)
+            return models

app/main.py

@@ -0,0 +1,350 @@
+from fastapi import FastAPI, HTTPException, Request, Depends, status
+from fastapi.responses import JSONResponse, StreamingResponse
+from .models import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse, ModelList
+from .gemini import GeminiClient, ResponseWrapper
+from .utils import handle_gemini_error, protect_from_abuse, APIKeyManager, test_api_key
+import os
+import json
+import asyncio
+from typing import Literal
+import random
+import requests
+from datetime import datetime, timedelta
+from apscheduler.schedulers.background import BackgroundScheduler
+import sys
+
+
+DEBUG = os.environ.get("DEBUG", "false").lower() == "true"
+LOG_FORMAT_DEBUG = '%(asctime)s - %(levelname)s - [%(key)s]-%(request_type)s-[%(model)s]-%(status_code)s: %(message)s - %(error_message)s'
+LOG_FORMAT_NORMAL = '[%(key)s]-%(request_type)s-[%(model)s]-%(status_code)s: %(message)s'
+
+
+def format_log_message(level, message, extra=None):
+    """格式化日志消息，模拟之前的 logging 格式"""
+    log_values = {
+        'asctime': datetime.now().strftime("%Y-%m-%d %H:%M:%S"), # 模拟 asctime
+        'levelname': level, # 日志级别
+        'key': extra.get('key', 'N/A') if extra else 'N/A',
+        'request_type': extra.get('request_type', 'N/A') if extra else 'N/A',
+        'model': extra.get('model', 'N/A') if extra else 'N/A',
+        'status_code': extra.get('status_code', 'N/A') if extra else 'N/A',
+        'error_message': extra.get('error_message', '') if extra else '' ,
+        'message': message
+    }
+    log_format = LOG_FORMAT_DEBUG if DEBUG else LOG_FORMAT_NORMAL
+    return log_format % log_values
+
+
+def translate_error(message: str) -> str:
+    if "quota exceeded" in message.lower():
+        return "API 密钥配额已用尽"
+    if "invalid argument" in message.lower():
+        return "无效参数"
+    if "internal server error" in message.lower():
+        return "服务器内部错误"
+    if "service unavailable" in message.lower():
+        return "服务不可用"
+    return message
+
+
+def handle_exception(exc_type, exc_value, exc_traceback):
+    if issubclass(exc_type, KeyboardInterrupt):
+        sys.excepthook(exc_type, exc_value, exc_traceback)
+        return
+    error_message = translate_error(str(exc_value))
+    log_msg = format_log_message('ERROR', f"未捕获的异常: %s" % error_message, extra={'status_code': 500, 'error_message': error_message})
+    print(log_msg)
+
+
+sys.excepthook = handle_exception
+
+app = FastAPI()
+
+PASSWORD = os.environ.get("PASSWORD", "123")
+MAX_REQUESTS_PER_MINUTE = int(os.environ.get("MAX_REQUESTS_PER_MINUTE", "30"))
+MAX_REQUESTS_PER_DAY_PER_IP = int(
+    os.environ.get("MAX_REQUESTS_PER_DAY_PER_IP", "600"))
+MAX_RETRIES = int(os.environ.get('MaxRetries', '3').strip() or '3')
+RETRY_DELAY = 1
+MAX_RETRY_DELAY = 16
+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"
+    },
+    {
+        "category": 'HARM_CATEGORY_CIVIC_INTEGRITY',
+        "threshold": 'BLOCK_NONE'
+    }
+]
+safety_settings_g2 = [
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "OFF"
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "OFF"
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "OFF"
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "OFF"
+    },
+    {
+        "category": 'HARM_CATEGORY_CIVIC_INTEGRITY',
+        "threshold": 'OFF'
+    }
+]
+
+key_manager = APIKeyManager() # 实例化 APIKeyManager，栈会在 __init__ 中初始化
+current_api_key = key_manager.get_available_key()
+
+
+def switch_api_key():
+    global current_api_key
+    key = key_manager.get_available_key() # get_available_key 会处理栈的逻辑
+    if key:
+        current_api_key = key
+        log_msg = format_log_message('INFO', f"API key 替换为 → {current_api_key[:8]}...", extra={'key': current_api_key[:8], 'request_type': 'switch_key'})
+        print(log_msg)
+    else:
+        log_msg = format_log_message('ERROR', "API key 替换失败，所有API key都已耗尽或被暂时禁用，请重新配置或稍后重试", extra={'key': 'N/A', 'request_type': 'switch_key', 'status_code': 'N/A'})
+        print(log_msg)
+
+
+async def check_keys():
+    available_keys = []
+    for key in key_manager.api_keys:
+        is_valid = await test_api_key(key)
+        status_msg = "有效" if is_valid else "无效"
+        log_msg = format_log_message('INFO', f"API Key {key[:10]}... {status_msg}.")
+        print(log_msg)
+        if is_valid:
+            available_keys.append(key)
+    if not available_keys:
+        log_msg = format_log_message('ERROR', "没有可用的 API 密钥！", extra={'key': 'N/A', 'request_type': 'startup', 'status_code': 'N/A'})
+        print(log_msg)
+    return available_keys
+
+
+@app.on_event("startup")
+async def startup_event():
+    log_msg = format_log_message('INFO', "Starting Gemini API proxy...")
+    print(log_msg)
+    available_keys = await check_keys()
+    if available_keys:
+        key_manager.api_keys = available_keys
+        key_manager._reset_key_stack() # 启动时也确保创建随机栈
+        key_manager.show_all_keys()
+        log_msg = format_log_message('INFO', f"可用 API 密钥数量：{len(key_manager.api_keys)}")
+        print(log_msg)
+        MAX_RETRIES = len(key_manager.api_keys)  # 动态设置 MAX_RETRIES 为密钥数量
+        log_msg = format_log_message('INFO', f"最大重试次数设置为：{MAX_RETRIES}") # 添加日志
+        print(log_msg)
+        if key_manager.api_keys:
+            all_models = await GeminiClient.list_available_models(key_manager.api_keys[0])
+            GeminiClient.AVAILABLE_MODELS = [model.replace(
+                "models/", "") for model in all_models]
+            log_msg = format_log_message('INFO', "Available models loaded.")
+            print(log_msg)
+
+@app.get("/v1/models", response_model=ModelList)
+def list_models():
+    log_msg = format_log_message('INFO', "Received request to list models", extra={'request_type': 'list_models', 'status_code': 200})
+    print(log_msg)
+    return ModelList(data=[{"id": model, "object": "model", "created": 1678888888, "owned_by": "organization-owner"} for model in GeminiClient.AVAILABLE_MODELS])
+
+
+async def verify_password(request: Request):
+    if PASSWORD:
+        auth_header = request.headers.get("Authorization")
+        if not auth_header or not auth_header.startswith("Bearer "):
+            raise HTTPException(
+                status_code=401, detail="Unauthorized: Missing or invalid token")
+        token = auth_header.split(" ")[1]
+        if token != PASSWORD:
+            raise HTTPException(
+                status_code=401, detail="Unauthorized: Invalid token")
+
+
+async def process_request(chat_request: ChatCompletionRequest, http_request: Request, request_type: Literal['stream', 'non-stream']):
+    global current_api_key
+    protect_from_abuse(
+        http_request, MAX_REQUESTS_PER_MINUTE, MAX_REQUESTS_PER_DAY_PER_IP)
+    if chat_request.model not in GeminiClient.AVAILABLE_MODELS:
+        error_msg = "无效的模型"
+        extra_log = {'request_type': request_type, 'model': chat_request.model, 'status_code': 400, 'error_message': error_msg}
+        log_msg = format_log_message('ERROR', error_msg, extra=extra_log)
+        print(log_msg)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail=error_msg)
+
+    key_manager.reset_tried_keys_for_request() # 在每次请求处理开始时重置 tried_keys 集合
+
+    contents, system_instruction = GeminiClient.convert_messages(
+        GeminiClient, chat_request.messages)
+
+    retry_attempts = len(key_manager.api_keys) if key_manager.api_keys else 1 # 重试次数等于密钥数量，至少尝试 1 次
+    for attempt in range(1, retry_attempts + 1):
+        extra_log_attempt_start = {'key': 'N/A', 'request_type': request_type, 'model': chat_request.model} # 初始 key 为 N/A
+        log_msg_attempt_start = format_log_message('INFO', f"第 {attempt}/{retry_attempts} 次尝试 ...", extra=extra_log_attempt_start)
+        print(log_msg_attempt_start)
+
+        current_api_key = key_manager.get_available_key() # 每次循环都获取新的 key, 栈逻辑在 get_available_key 中处理
+
+        if current_api_key is None: # 检查是否获取到 API 密钥
+            log_msg_no_key = format_log_message('WARNING', "没有可用的 API 密钥，跳过本次尝试", extra={'request_type': request_type, 'model': chat_request.model, 'status_code': 'N/A'})
+            print(log_msg_no_key)
+            break  # 如果没有可用密钥，跳出循环
+
+        extra_log = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model} # 使用获取到的 key 记录日志
+        log_msg = format_log_message('INFO', f"第 {attempt}/{retry_attempts} 次尝试 ... 使用密钥: {current_api_key[:8]}...", extra=extra_log)
+        print(log_msg)
+
+
+        gemini_client = GeminiClient(current_api_key)
+        try:
+            if chat_request.stream:
+                async def stream_generator():
+                    try:
+                        async for chunk in gemini_client.stream_chat(chat_request, contents, safety_settings_g2 if 'gemini-2.0-flash-exp' in chat_request.model else safety_settings, system_instruction):
+                            formatted_chunk = {"id": "chatcmpl-someid", "object": "chat.completion.chunk", "created": 1234567,
+                                               "model": chat_request.model, "choices": [{"delta": {"role": "assistant", "content": chunk}, "index": 0, "finish_reason": None}]}
+                            yield f"data: {json.dumps(formatted_chunk)}\n\n"
+                        yield "data: [DONE]\n\n"
+
+                    except asyncio.CancelledError:
+                        extra_log_cancel = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'error_message': '客户端已断开连接'}
+                        log_msg = format_log_message('INFO', "Client disconnected", extra=extra_log_cancel)
+                        print(log_msg)
+                    except Exception as e:
+                        error_detail = handle_gemini_error(
+                            e, current_api_key, key_manager, switch_api_key)
+                        log_message = f"API Key failed: {error_detail}"
+                        extra_log_error = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'status_code': 500, 'error_message': error_detail}
+                        log_msg = format_log_message('ERROR', log_message, extra=extra_log_error)
+                        print(log_msg)
+                        yield f"data: {json.dumps({'error': {'message': error_detail, 'type': 'gemini_error'}})}\n\n"
+                        if attempt < retry_attempts: # 流式也根据apikey 数量判断是否切换key
+                            switch_api_key() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+                return StreamingResponse(stream_generator(), media_type="text/event-stream")
+            else:
+                async def run_gemini_completion():
+                    try:
+                        response_content = await asyncio.to_thread(gemini_client.complete_chat, chat_request, contents, safety_settings_g2 if 'gemini-2.0-flash-exp' in chat_request.model else safety_settings, system_instruction)
+                        return response_content
+                    except asyncio.CancelledError:
+                        extra_log_gemini_cancel = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'error_message': 'Gemini API 调用因客户端断开连接而被取消'}
+                        log_msg = format_log_message('INFO', "Gemini API call cancelled due to client disconnect", extra=extra_log_gemini_cancel)
+                        print(log_msg)
+                        raise
+
+                async def check_client_disconnect():
+                    while True:
+                        if await http_request.is_disconnected():
+                            extra_log_client_disconnect = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'error_message': '在非流式请求期间检测到客户端断开连接。正在取消 Gemini API 调用。'}
+                            log_msg = format_log_message('INFO', "Client disconnected during non-streaming request.  Cancelling Gemini API call.", extra=extra_log_client_disconnect)
+                            print(log_msg)
+                            return True
+                        await asyncio.sleep(0.5)
+
+                gemini_task = asyncio.create_task(run_gemini_completion())
+                disconnect_task = asyncio.create_task(check_client_disconnect())
+
+                try:
+                    done, pending = await asyncio.wait(
+                        [gemini_task, disconnect_task],
+                        return_when=asyncio.FIRST_COMPLETED
+                    )
+
+                    if disconnect_task in done:
+                        gemini_task.cancel()
+                        try:
+                            await gemini_task
+                        except asyncio.CancelledError:
+                            extra_log_gemini_task_cancel = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'error_message': '客户端断开连接后，Gemini API 任务已成功取消。'}
+                            log_msg = format_log_message('INFO', "Gemini API task successfully cancelled after client disconnect.", extra=extra_log_gemini_task_cancel)
+                            print(log_msg)
+                            pass
+                        raise HTTPException(status_code=status.HTTP_408_REQUEST_TIMEOUT, detail="Client disconnected")
+
+                    if gemini_task in done:
+                        disconnect_task.cancel()
+                        try:
+                            await disconnect_task
+                        except asyncio.CancelledError:
+                            pass
+                        response_content = gemini_task.result()
+                        response = ChatCompletionResponse(id="chatcmpl-someid", object="chat.completion", created=1234567890, model=chat_request.model,
+                                                        choices=[{"index": 0, "message": {"role": "assistant", "content": response_content.text}, "finish_reason": "stop"}])
+                        extra_log_success = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'status_code': 200}
+                        log_msg = format_log_message('INFO', "Request successful", extra=extra_log_success)
+                        print(log_msg)
+                        return response
+
+                except asyncio.CancelledError:
+                    extra_log_request_cancel = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'error_message':"请求被取消" }
+                    log_msg = format_log_message('INFO', "Request cancelled", extra=extra_log_request_cancel)
+                    print(log_msg)
+                    raise
+
+
+        except requests.exceptions.RequestException as e:
+            error_detail = handle_gemini_error(
+                e, current_api_key, key_manager, switch_api_key)
+            extra_log_request_exception = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'status_code': 500, 'error_message': error_detail}
+            log_msg = format_log_message('ERROR', f"{error_detail}", extra=extra_log_request_exception)
+            print(log_msg)
+            if attempt < retry_attempts: # 根据apikey 数量判断是否切换key
+                switch_api_key() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"{retry_attempts} 次尝试后仍然失败，请修改预设或输入") # 错误信息里的重试次数也动态修改
+        except Exception as e:
+            error_detail = handle_gemini_error(
+                e, current_api_key, key_manager, switch_api_key)
+            extra_log_exception = {'key': current_api_key[:8], 'request_type': request_type, 'model': chat_request.model, 'status_code': 500, 'error_message': error_detail}
+            log_msg = format_log_message('ERROR', f"{error_detail}", extra=extra_log_exception)
+            print(log_msg)
+            if attempt < retry_attempts: # 根据apikey 数量判断是否切换key
+                switch_api_key() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"{retry_attempts} 次尝试后仍然失败，请修改预设或输入") # 错误信息里的重试次数也动态修改
+
+    msg = "所有API密钥或重试次数均失败"
+    extra_log_all_fail = {'key': "ALL", 'request_type': request_type, 'model': chat_request.model, 'status_code': 500, 'error_message': msg}
+    log_msg = format_log_message('ERROR', msg, extra=extra_log_all_fail)
+    print(log_msg)
+    raise HTTPException(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=msg)
+
+
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(request: ChatCompletionRequest, http_request: Request, _: None = Depends(verify_password)):
+    return await process_request(request, http_request, "stream" if request.stream else "non-stream")
+
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    error_message = translate_error(str(exc))
+    extra_log_unhandled_exception = {'status_code': 500, 'error_message': error_message}
+    log_msg = format_log_message('ERROR', f"Unhandled exception: {error_message}", extra=extra_log_unhandled_exception)
+    print(log_msg)
+    return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=ErrorResponse(message=str(exc), type="internal_error").dict())

app/models.py

@@ -0,0 +1,46 @@
+from typing import List, Dict, Optional, Union, Literal
+from pydantic import BaseModel, Field
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    temperature: float = 0.7
+    top_p: Optional[float] = 1.0
+    n: int = 1
+    stream: bool = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+
+class Choice(BaseModel):
+    index: int
+    message: Message
+    finish_reason: Optional[str] = None
+
+class Usage(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: Literal["chat.completion"]
+    created: int
+    model: str
+    choices: List[Choice]
+    usage: Usage = Field(default_factory=Usage)
+
+class ErrorResponse(BaseModel):
+    message: str
+    type: str
+    param: Optional[str] = None
+    code: Optional[str] = None
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data: List[Dict]

app/utils.py

@@ -0,0 +1,232 @@
+import random
+from fastapi import HTTPException, Request
+import time
+import re
+from datetime import datetime, timedelta
+from apscheduler.schedulers.background import BackgroundScheduler
+import os
+import requests
+import httpx
+from threading import Lock
+
+DEBUG = os.environ.get("DEBUG", "false").lower() == "true"
+LOG_FORMAT_DEBUG = '%(asctime)s - %(levelname)s - [%(key)s]-%(request_type)s-[%(model)s]-%(status_code)s: %(message)s - %(error_message)s'
+LOG_FORMAT_NORMAL = '[%(key)s]-%(request_type)s-[%(model)s]-%(status_code)s: %(message)s'
+
+
+def format_log_message(level, message, extra=None):
+    """格式化日志消息，模拟之前的 logging 格式"""
+    log_values = {
+        'asctime': datetime.now().strftime("%Y-%m-%d %H:%M:%S"), # 模拟 asctime
+        'levelname': level, # 日志级别
+        'key': extra.get('key', 'N/A') if extra else 'N/A',
+        'request_type': extra.get('request_type', 'N/A') if extra else 'N/A',
+        'model': extra.get('model', 'N/A') if extra else 'N/A',
+        'status_code': extra.get('status_code', 'N/A') if extra else 'N/A',
+        'error_message': extra.get('error_message', '') if extra else '' ,
+        'message': message
+    }
+    log_format = LOG_FORMAT_DEBUG if DEBUG else LOG_FORMAT_NORMAL
+    return log_format % log_values
+
+
+class APIKeyManager:
+    def __init__(self):
+        self.api_keys = re.findall(
+            r"AIzaSy[a-zA-Z0-9_-]{33}", os.environ.get('GEMINI_API_KEYS', ""))
+        self.key_stack = [] # 初始化密钥栈
+        self._reset_key_stack() # 初始化时创建随机密钥栈
+        self.api_key_blacklist = set()
+        self.api_key_blacklist_duration = 60
+        self.scheduler = BackgroundScheduler()
+        self.scheduler.start()
+        self.tried_keys_for_request = set()  # 用于跟踪当前请求尝试中已试过的 key
+
+    def _reset_key_stack(self):
+        """创建并随机化密钥栈"""
+        shuffled_keys = self.api_keys[:]  # 创建 api_keys 的副本以避免直接修改原列表
+        random.shuffle(shuffled_keys)
+        self.key_stack = shuffled_keys
+        log_msg = format_log_message('INFO', "已重新创建随机密钥栈", extra={'request_type': 'key_stack', 'status_code': 'N/A'})
+        print(log_msg)
+
+
+    def get_available_key(self):
+        """从栈顶获取密钥，栈空时重新生成 (修改后)"""
+        while self.key_stack:
+            key = self.key_stack.pop()
+            if key not in self.api_key_blacklist and key not in self.tried_keys_for_request:
+                self.tried_keys_for_request.add(key)
+                return key
+
+        if not self.api_keys:
+            log_msg = format_log_message('ERROR', "没有配置任何 API 密钥！")
+            print(log_msg)
+            return None
+
+        log_msg = format_log_message('WARNING', "密钥栈已用尽或栈内密钥均不可用，重新生成密钥栈")
+        print(log_msg)
+        self._reset_key_stack() # 重新生成密钥栈
+
+        # 再次尝试从新栈中获取密钥 (迭代一次)
+        while self.key_stack:
+            key = self.key_stack.pop()
+            if key not in self.api_key_blacklist and key not in self.tried_keys_for_request:
+                self.tried_keys_for_request.add(key)
+                return key
+
+        return None
+
+
+    def show_all_keys(self):
+        log_msg = format_log_message('INFO', f"当前可用API key个数: {len(self.api_keys)} ")
+        print(log_msg)
+        for i, api_key in enumerate(self.api_keys):
+            log_msg = format_log_message('INFO', f"API Key{i}: {api_key[:8]}...{api_key[-3:]}")
+            print(log_msg)
+
+    def blacklist_key(self, key):
+        log_msg = format_log_message('WARNING', f"{key[:8]} → 暂时禁用 {self.api_key_blacklist_duration} 秒")
+        print(log_msg)
+        self.api_key_blacklist.add(key)
+        self.scheduler.add_job(lambda: self.api_key_blacklist.discard(key), 'date',
+                               run_date=datetime.now() + timedelta(seconds=self.api_key_blacklist_duration))
+
+    def reset_tried_keys_for_request(self):
+        """在新的请求尝试时重置已尝试的 key 集合"""
+        self.tried_keys_for_request = set()
+
+
+def handle_gemini_error(error, current_api_key, key_manager, switch_api_key_func) -> str:
+    if isinstance(error, requests.exceptions.HTTPError):
+        status_code = error.response.status_code
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                if 'error' in error_data:
+                    if error_data['error'].get('code') == "invalid_argument":
+                        error_message = "无效的 API 密钥"
+                        extra_log_invalid_key = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+                        log_msg = format_log_message('ERROR', f"{current_api_key[:8]} ... {current_api_key[-3:]} → 无效，可能已过期或被删除", extra=extra_log_invalid_key)
+                        print(log_msg)
+                        key_manager.blacklist_key(current_api_key)
+                        switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+                        return error_message
+                    error_message = error_data['error'].get(
+                        'message', 'Bad Request')
+                    extra_log_400 = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+                    log_msg = format_log_message('WARNING', f"400 错误请求: {error_message}", extra=extra_log_400)
+                    print(log_msg)
+                    return f"400 错误请求: {error_message}"
+            except ValueError:
+                error_message = "400 错误请求：响应不是有效的JSON格式"
+                extra_log_400_json = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+                log_msg = format_log_message('WARNING', error_message, extra=extra_log_400_json)
+                print(log_msg)
+                return error_message
+
+        elif status_code == 429:
+            error_message = "API 密钥配额已用尽"
+            extra_log_429 = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+            log_msg = format_log_message('WARNING', f"{current_api_key[:8]} ... {current_api_key[-3:]} → 429 官方资源耗尽", extra=extra_log_429)
+            print(log_msg)
+            key_manager.blacklist_key(current_api_key)
+            switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            return error_message
+
+        elif status_code == 403:
+            error_message = "权限被拒绝"
+            extra_log_403 = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+            log_msg = format_log_message('ERROR', f"{current_api_key[:8]} ... {current_api_key[-3:]} → 403 权限被拒绝", extra=extra_log_403)
+            print(log_msg)
+            key_manager.blacklist_key(current_api_key)
+            switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            return error_message
+        elif status_code == 500:
+            error_message = "服务器内部错误"
+            extra_log_500 = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+            log_msg = format_log_message('WARNING', f"{current_api_key[:8]} ... {current_api_key[-3:]} → 500 服务器内部错误", extra=extra_log_500)
+            print(log_msg)
+            switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            return "Gemini API 内部错误"
+
+        elif status_code == 503:
+            error_message = "服务不可用"
+            extra_log_503 = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+            log_msg = format_log_message('WARNING', f"{current_api_key[:8]} ... {current_api_key[-3:]} → 503 服务不可用", extra=extra_log_503)
+            print(log_msg)
+            switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            return "Gemini API 服务不可用"
+        else:
+            error_message = f"未知错误: {status_code}"
+            extra_log_other = {'key': current_api_key[:8], 'status_code': status_code, 'error_message': error_message}
+            log_msg = format_log_message('WARNING', f"{current_api_key[:8]} ... {current_api_key[-3:]} → {status_code} 未知错误", extra=extra_log_other)
+            print(log_msg)
+            switch_api_key_func() # 这里虽然叫 switch_api_key_func, 但实际上 get_available_key 会处理栈和重新生成
+            return f"未知错误/模型不可用: {status_code}"
+
+    elif isinstance(error, requests.exceptions.ConnectionError):
+        error_message = "连接错误"
+        log_msg = format_log_message('WARNING', error_message, extra={'error_message': error_message})
+        print(log_msg)
+        return error_message
+
+    elif isinstance(error, requests.exceptions.Timeout):
+        error_message = "请求超时"
+        log_msg = format_log_message('WARNING', error_message, extra={'error_message': error_message})
+        print(log_msg)
+        return error_message
+    else:
+        error_message = f"发生未知错误: {error}"
+        log_msg = format_log_message('ERROR', error_message, extra={'error_message': error_message})
+        print(log_msg)
+        return error_message
+
+
+async def test_api_key(api_key: str) -> bool:
+    """
+    测试 API 密钥是否有效���
+    """
+    try:
+        url = "https://generativelanguage.googleapis.com/v1beta/models?key={}".format(api_key)
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url)
+            response.raise_for_status()
+            return True
+    except Exception:
+        return False
+
+
+rate_limit_data = {}
+rate_limit_lock = Lock()
+
+
+def protect_from_abuse(request: Request, max_requests_per_minute: int = 30, max_requests_per_day_per_ip: int = 600):
+    now = int(time.time())
+    minute = now // 60
+    day = now // (60 * 60 * 24)
+
+    minute_key = f"{request.url.path}:{minute}"
+    day_key = f"{request.client.host}:{day}"
+
+    with rate_limit_lock:
+        minute_count, minute_timestamp = rate_limit_data.get(
+            minute_key, (0, now))
+        if now - minute_timestamp >= 60:
+            minute_count = 0
+            minute_timestamp = now
+        minute_count += 1
+        rate_limit_data[minute_key] = (minute_count, minute_timestamp)
+
+        day_count, day_timestamp = rate_limit_data.get(day_key, (0, now))
+        if now - day_timestamp >= 86400:
+            day_count = 0
+            day_timestamp = now
+        day_count += 1
+        rate_limit_data[day_key] = (day_count, day_timestamp)
+
+    if minute_count > max_requests_per_minute:
+        raise HTTPException(status_code=429, detail={
+            "message": "Too many requests per minute", "limit": max_requests_per_minute})
+    if day_count > max_requests_per_day_per_ip:
+        raise HTTPException(status_code=429, detail={"message": "Too many requests per day from this IP", "limit": max_requests_per_day_per_ip})