app.py

import json
import re
import time
import asyncio
import uvicorn
from fastapi import FastAPI, Request, HTTPException, Header, Depends
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Union
import requests
from datetime import datetime
import logging
import os
from dotenv import load_dotenv

from apscheduler.schedulers.background import BackgroundScheduler

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("openai-proxy")

# 创建FastAPI应用
app = FastAPI(
    title="OpenAI API Proxy",
    description="将OpenAI API请求代理到DeepSider API",
    version="1.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置
DEEPSIDER_API_BASE = "https://api.chargpt.ai/api/v2"
DEEPSIDER_TOKEN = os.getenv("DEEPSIDER_TOKEN", "").split(',')
TOKEN_INDEX = 0

# 模型映射表
MODEL_MAPPING = {
    "gpt-4o-mini": "openai/gpt-4o-mini",
    "gpt-4o": "openai/gpt-4o",
    "o1": "openai/o1",
    "o3-mini": "openai/o3-mini",
    "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
    "claude-3.7-sonnet": "anthropic/claude-3.7-sonnet",
    "grok-3": "x-ai/grok-3",
    "grok-3-reasoner": "x-ai/grok-3-reasoner",
    "deepseek-v3": "deepseek/deepseek-chat",
    "deepseek-r1": "deepseek/deepseek-r1",
    "gemini-2.0-flash": "google/gemini-2.0-flash",
    "gemini-2.0-pro-exp": "google/gemini-2.0-pro-exp-02-05",
    "gemini-2.0-flash-thinking-exp": "google/gemini-2.0-flash-thinking-exp-1219",
    "qwq-32b": "qwen/qwq-32b",
    "qwen-max": "qwen/qwen-max",
    "DeepSeek-V3-0324": "deepseek/deepseek-chat-v3-0324",
    "gpt-4o-image": "openai/gpt-4o-image"
}

# Token负载均衡状态
token_status = {}


# 请求头
def get_headers():
    global TOKEN_INDEX
    # 负载均衡，轮询选择token
    if len(DEEPSIDER_TOKEN) > 0:
        current_token = DEEPSIDER_TOKEN[TOKEN_INDEX % len(DEEPSIDER_TOKEN)]
        # TOKEN_INDEX = (TOKEN_INDEX + 1) % len(DEEPSIDER_TOKEN)

        # 检查token状态
        if current_token in token_status and not token_status[current_token]["active"]:
            # 如果token不可用，尝试下一个
            for i in range(len(DEEPSIDER_TOKEN)):
                next_token = DEEPSIDER_TOKEN[(TOKEN_INDEX + i) % len(DEEPSIDER_TOKEN)]
                if next_token not in token_status or token_status[next_token]["active"]:
                    current_token = next_token
                    TOKEN_INDEX = (TOKEN_INDEX + i + 1) % len(DEEPSIDER_TOKEN)
                    break
    else:
        current_token = ""

    return {
        "accept": "*/*",
        "accept-encoding": "gzip, deflate, br, zstd",
        "accept-language": "en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7",
        "content-type": "application/json",
        "origin": "chrome-extension://client",
        "i-lang": "zh-CN",
        "i-version": "1.1.64",
        "sec-ch-ua": '"Chromium";v="134", "Not:A-Brand";v="24"',
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": "Windows",
        "sec-fetch-dest": "empty",
        "sec-fetch-mode": "cors",
        "sec-fetch-site": "cross-site",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36",
        "authorization": f"Bearer {current_token}"
    }


# OpenAI API请求模型
class ChatMessage(BaseModel):
    role: str
    content: str
    name: Optional[str] = None


class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 1.0
    top_p: Optional[float] = 1.0
    n: Optional[int] = 1
    stream: Optional[bool] = False
    stop: Optional[Union[List[str], str]] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = 0
    frequency_penalty: Optional[float] = 0
    user: Optional[str] = None


def reset_task():
    try:
        for key, value in token_status.items():
            token_status[key]["total"] = value["count"]
            token_status[key]["active"] = True
        print(f"执行重置任务... 当前时间: {datetime.now()}")
    except Exception as e:
        print(f"任务执行出错: {e}")


# 初始化token状态
async def initialize_token_status():
    """初始化检查所有token的状态和余额"""
    global token_status

    for token in DEEPSIDER_TOKEN:
        headers = {
            "accept": "*/*",
            "content-type": "application/json",
            "authorization": f"Bearer {token}"
        }

        try:
            # 获取账户余额信息
            response = requests.get(
                f"{DEEPSIDER_API_BASE.replace('/v2', '')}/quota/retrieve",
                headers=headers
            )

            active = False
            quota_info = {}
            count = 0
            total = 0

            if response.status_code == 200:
                data = response.json()
                if data.get('code') == 0:
                    quota_list = data.get('data', {}).get('list', [])

                    # 解析余额信息
                    for item in quota_list:
                        item_type = item.get('type', '')
                        available = item.get('available', 0)
                        count += available
                        total += item.get('total', 0)
                        if available > 0:
                            active = True

                        quota_info[item_type] = {
                            "total": item.get('total', 0),
                            "available": available,
                            "title": item.get('title', '')
                        }

            token_status[token] = {
                "active": active,
                "quota": quota_info,
                "last_checked": datetime.now(),
                "failed_count": 0,
                "count": count,
                'total': total
            }

            logger.info(f"Token {token[:8]}... 状态：{'活跃' if active else '无效'}")

        except Exception as e:
            logger.warning(f"检查Token {token[:8]}... 出错：{str(e)}")
            token_status[token] = {
                "active": False,
                "quota": {},
                "last_checked": datetime.now(),
                "failed_count": 0
            }


# 工具函数
def verify_api_key(api_key: str = Header(..., alias="Authorization")):
    """验证API密钥"""
    if not api_key.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Invalid API key format")
    return api_key.replace("Bearer ", "")


def map_openai_to_deepsider_model(model: str) -> str:
    """将OpenAI模型名称映射到DeepSider模型名称"""
    return MODEL_MAPPING.get(model, model)


def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
    """格式化消息列表为DeepSider API所需的提示格式"""
    prompt = ""
    next_code = False
    for msg in messages:

        if next_code:
            next_code = False
            continue

        if msg.role == 'assistant' and (
                '验证码提示' in msg.content or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
            next_code = True
            continue

        role = msg.role
        # 将OpenAI的角色映射到DeepSider能理解的格式
        if role == "system":
            # 系统消息放在开头 作为指导
            prompt = f"{msg.content}\n\n" + prompt
        elif role == "user":
            prompt += f"Human: {msg.content}\n\n"
        elif role == "assistant":
            content = msg.content
            re.sub(r'\[clId:(.*)]', '', content)
            prompt += f"Assistant: {msg.content}\n\n"

        else:
            # 其他角色按用户处理
            prompt += f"Human ({role}): {msg.content}\n\n"

    # 如果最后一个消息不是用户的 添加一个Human前缀引导模型回答
    if messages and messages[-1].role != "user":
        prompt += "Human: "

    return prompt.strip()


def update_token_status(token: str, success: bool, error_message: str = None):
    """更新token的状态"""
    global token_status

    if token not in token_status:
        token_status[token] = {
            "active": True,
            "quota": {},
            "last_checked": datetime.now(),
            "failed_count": 0
        }

    if not success:
        token_status[token]["failed_count"] += 1

        # 如果失败消息包含余额不足，标记为不活跃
        if error_message and ("配额不足" in error_message or "quota" in error_message.lower()):
            token_status[token]["active"] = False
            logger.warning(f"Token {token[:8]}... 余额不足，已标记为不活跃")

        # 连续失败5次，也标记为不活跃
        if token_status[token]["failed_count"] >= 5:
            token_status[token]["active"] = False
            logger.warning(f"Token {token[:8]}... 连续失败{token_status[token]['failed_count']}次，已标记为不活跃")
    else:
        # 成功则重置失败计数
        token_status[token]["count"] -= 1
        token_status[token]["failed_count"] = 0

    if token_status[token]["count"] <= 0:
        token_status[token]["active"] = False
        logger.warning(f"Token {token[:8]}... 余额不足，已标记为不活跃")


async def generate_openai_response(full_response: str, request_id: str, model: str) -> Dict:
    """生成符合OpenAI API响应格式的完整响应"""
    timestamp = int(time.time())
    return {
        "id": f"chatcmpl-{request_id}",
        "object": "chat.completion",
        "created": timestamp,
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": full_response
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": 0,  # 无法准确计算
            "completion_tokens": 0,  # 无法准确计算
            "total_tokens": 0  # 无法准确计算
        }
    }


async def stream_openai_response(response, request_id: str, model: str, token: str):
    global TOKEN_INDEX
    """流式返回OpenAI API格式的响应"""
    timestamp = int(time.time())
    full_response = ""
    codeFlag = False

    try:
        # 将DeepSider响应流转换为OpenAI流格式
        for line in response.iter_lines():
            if not line:
                continue

            if line.startswith(b'data: '):
                try:
                    data = json.loads(line[6:].decode('utf-8'))
                    if data.get('code') == 1005:
                        raise Exception(data.get("message"))

                    if data.get('code') == 202 and data.get('data', {}).get('type') == "chat":
                        # 获取正文内容
                        content = data.get('data', {}).get('content', '')
                        if content:
                            full_response += content

                            # 生成OpenAI格式的流式响应
                            chunk = {
                                "id": f"chatcmpl-{request_id}",
                                "object": "chat.completion.chunk",
                                "created": timestamp,
                                "model": model,
                                "choices": [
                                    {
                                        "index": 0,
                                        "delta": {
                                            "content": content
                                        },
                                        "finish_reason": None
                                    }
                                ]
                            }

                            if '验证码提示' in content or '验证码已发送至您的邮箱' in content:
                                codeFlag = True
                            yield f"data: {json.dumps(chunk)}\n\n"

                    elif data.get('code') == 203:
                        # 尾巴
                        if codeFlag or model == 'gpt-4o-image':
                            chunk = {
                                "id": f"chatcmpl-{request_id}",
                                "object": "chat.completion.chunk",
                                "created": timestamp,
                                "model": model,
                                "choices": [
                                    {
                                        "index": 0,
                                        "delta": {
                                            "content": f"\n[clId:{data.get('data', {}).get('clId')}]"
                                        },
                                        "finish_reason": None
                                    }
                                ]
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"

                        # 生成完成信号
                        chunk = {
                            "id": f"chatcmpl-{request_id}",
                            "object": "chat.completion.chunk",
                            "created": timestamp,
                            "model": model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {},
                                    "finish_reason": "stop"
                                }
                            ]
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
                        yield "data: [DONE]\n\n"

                except json.JSONDecodeError:
                    logger.warning(f"无法解析响应: {line}")

        # 更新token状态（成功）
        if not codeFlag:
            update_token_status(token, True)

    except Exception as e:
        logger.error(f"流式响应处理出错: {str(e)}")
        if '今日额度已用完啦' in str(e):
            TOKEN_INDEX = (TOKEN_INDEX + 1) % len(DEEPSIDER_TOKEN)
        # 更新token状态（失败）
        update_token_status(token, False, str(e))

        # 返回错误信息
        error_chunk = {
            "id": f"chatcmpl-{request_id}",
            "object": "chat.completion.chunk",
            "created": timestamp,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {
                        "content": f"\n\n[处理响应时出错: {str(e)}]"
                    },
                    "finish_reason": "stop"
                }
            ]
        }
        yield f"data: {json.dumps(error_chunk)}\n\n"
        yield "data: [DONE]\n\n"


# 路由定义
@app.get("/")
async def root():
    return {"message": "OpenAI API Proxy服务已启动 连接至DeepSider API"}


@app.get("/v1/models")
async def list_models(api_key: str = Depends(verify_api_key)):
    """列出可用的模型"""
    models = []
    for openai_model, _ in MODEL_MAPPING.items():
        models.append({
            "id": openai_model,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "openai-proxy"
        })

    return {
        "object": "list",
        "data": models
    }


@app.post("/v1/chat/completions")
async def create_chat_completion(
        request: Request,
        api_key: str = Depends(verify_api_key)
):
    """创建聊天完成API - 支持普通请求和流式请求"""
    # 解析请求体
    body = await request.json()
    chat_request = ChatCompletionRequest(**body)

    # 生成唯一请求ID
    request_id = datetime.now().strftime("%Y%m%d%H%M%S") + str(time.time_ns())[-6:]

    # 映射模型
    deepsider_model = map_openai_to_deepsider_model(chat_request.model)
    model = chat_request.model
    # 验证验证码
    isCode = False
    clId = None

    if len(chat_request.messages) > 1:
        msg = chat_request.messages[-2]
        if msg.role == 'assistant' and (
                '验证码提示' in msg.content or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
            isCode = True
        
        pattern = r'\[clId:(.*)]'
        match = re.search(pattern, msg.content)
        if match:
            clId = match.group(1)

        # 准备DeepSider API所需的提示
    if isCode or model == 'gpt-4o-image':
        prompt = chat_request.messages[-1].content
    else:
        prompt = format_messages_for_deepsider(chat_request.messages)

    # 准备请求体
    payload = {
        "model": deepsider_model,
        "prompt": prompt,
        "webAccess": "close",  # 默认关闭网络访问
        "timezone": "Asia/Shanghai"
    }

    if isCode or model == 'gpt-4o-image':
        payload["clId"] = clId

    # 获取当前token
    headers = get_headers()
    current_token = headers["authorization"].replace("Bearer ", "")

    try:
        # 发送请求到DeepSider API
        response = requests.post(
            f"{DEEPSIDER_API_BASE}/chat/conversation",
            headers=headers,
            json=payload,
            stream=True
        )

        # 检查响应状态
        if response.status_code != 200:
            error_msg = f"DeepSider API请求失败: {response.status_code}"
            try:
                error_data = response.json()
                error_msg += f" - {error_data.get('message', '')}"
            except:
                error_msg += f" - {response.text}"

            logger.error(error_msg)

            # 更新token状态
            update_token_status(current_token, False, error_msg)

            raise HTTPException(status_code=response.status_code, detail="API请求失败")

        # 处理流式或非流式响应
        if chat_request.stream:
            # 返回流式响应
            return StreamingResponse(
                stream_openai_response(response, request_id, chat_request.model, current_token),
                media_type="text/event-stream"
            )
        else:
            # 收集完整响应
            full_response = ""
            for line in response.iter_lines():
                if not line:
                    continue

                if line.startswith(b'data: '):
                    try:
                        data = json.loads(line[6:].decode('utf-8'))

                        if data.get('code') == 202 and data.get('data', {}).get('type') == "chat":
                            content = data.get('data', {}).get('content', '')
                            if content:
                                full_response += content

                    except json.JSONDecodeError:
                        pass

            # 更新token状态（成功）
            update_token_status(current_token, True)

            # 返回OpenAI格式的完整响应
            return await generate_openai_response(full_response, request_id, chat_request.model)

    except HTTPException:
        raise
    except Exception as e:
        logger.exception("处理请求时出错")
        # 更新token状态（失败）
        update_token_status(current_token, False, str(e))
        raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")


# 查看token状态的端点
@app.get("/admin/tokens")
async def get_token_status(admin_key: str = Header(None, alias="X-Admin-Key")):
    """查看所有token的状态"""
    # 简单的管理密钥检查
    expected_admin_key = os.getenv("ADMIN_KEY", "admin")
    if not admin_key or admin_key != expected_admin_key:
        raise HTTPException(status_code=403, detail="Unauthorized")

    # 脱敏token，只显示前8位
    safe_status = {}
    for token, status in token_status.items():
        token_display = token[:8] + "..." if len(token) > 8 else token
        safe_status[token_display] = status

    return {"tokens": safe_status, "active_tokens": sum(1 for s in token_status.values() if s["active"])}


# 手动刷新token状态
@app.post("/admin/refresh-tokens")
async def refresh_token_status(admin_key: str = Header(None, alias="X-Admin-Key")):
    """手动刷新所有token的状态"""
    # 简单的管理密钥检查
    expected_admin_key = os.getenv("ADMIN_KEY", "admin")
    if not admin_key or admin_key != expected_admin_key:
        raise HTTPException(status_code=403, detail="Unauthorized")

    await ()
    return {"message": "所有token状态已刷新", "active_tokens": sum(1 for s in token_status.values() if s["active"])}


# 模拟模型的路由
@app.get("/v1/engines")
@app.get("/v1/engines/{engine_id}")
async def engines_handler():
    """兼容旧的引擎API"""
    raise HTTPException(status_code=404, detail="引擎API已被弃用 请使用模型API")


# 错误处理器
@app.exception_handler(404)
async def not_found_handler(request, exc):
    return {
        "error": {
            "message": f"未找到资源: {request.url.path}",
            "type": "not_found_error",
            "code": "not_found"
        }
    }, 404


# 启动事件
@app.on_event("startup")
async def startup_event():
    """服务启动时初始化token状态"""
    if not DEEPSIDER_TOKEN or (len(DEEPSIDER_TOKEN) == 1 and DEEPSIDER_TOKEN[0] == ""):
        logger.warning("未设置DEEPSIDER_TOKEN环境变量 请设置后再重启服务")
    else:
        logger.info(f"初始化 {len(DEEPSIDER_TOKEN)} 个token状态...")
        await initialize_token_status()
        active_tokens = sum(1 for s in token_status.values() if s["active"])
        logger.info(f"初始化完成 活跃token: {active_tokens}/{len(DEEPSIDER_TOKEN)}")


# 主程序
if __name__ == "__main__":
    scheduler = BackgroundScheduler()

    # 添加任务，每天0点执行
    scheduler.add_job(reset_task, 'cron', hour=0, minute=0)
    # 启动调度器
    scheduler.start()

    # 启动服务器
    port = int(os.getenv("PORT", "3000"))
    logger.info(f"启动OpenAI API代理服务 端口: {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)