app/api/routes.py

import json
from typing import Optional, Union
from fastapi import APIRouter, Body, HTTPException, Path, Query, Request, Depends, status, Header
from fastapi.responses import StreamingResponse
from app.services import GeminiClient
from app.utils import protect_from_abuse,generate_cache_key,openAI_from_text,log
from app.utils.response import openAI_from_Gemini
from app.utils.auth import custom_verify_password
from .stream_handlers import process_stream_request
from .nonstream_handlers import process_request, process_nonstream_with_keepalive_stream
from app.models.schemas import ChatCompletionRequest, ChatCompletionResponse, ModelList, AIRequest, ChatRequestGemini
import app.config.settings as settings
import asyncio
from app.vertex.routes import chat_api, models_api
from app.vertex.models import OpenAIRequest, OpenAIMessage

# 创建路由器
router = APIRouter()

# 全局变量引用 - 这些将在main.py中初始化并传递给路由
key_manager = None
response_cache_manager = None
active_requests_manager = None
safety_settings = None
safety_settings_g2 = None
current_api_key = None
FAKE_STREAMING = None
FAKE_STREAMING_INTERVAL = None
PASSWORD = None
MAX_REQUESTS_PER_MINUTE = None
MAX_REQUESTS_PER_DAY_PER_IP = None

# 初始化路由器的函数
def init_router(
    _key_manager,
    _response_cache_manager,
    _active_requests_manager,
    _safety_settings,
    _safety_settings_g2,
    _current_api_key,
    _fake_streaming,
    _fake_streaming_interval,
    _password,
    _max_requests_per_minute,
    _max_requests_per_day_per_ip
):
    global key_manager, response_cache_manager, active_requests_manager
    global safety_settings, safety_settings_g2, current_api_key
    global FAKE_STREAMING, FAKE_STREAMING_INTERVAL
    global PASSWORD, MAX_REQUESTS_PER_MINUTE, MAX_REQUESTS_PER_DAY_PER_IP
    
    key_manager = _key_manager
    response_cache_manager = _response_cache_manager
    active_requests_manager = _active_requests_manager
    safety_settings = _safety_settings
    safety_settings_g2 = _safety_settings_g2
    current_api_key = _current_api_key
    FAKE_STREAMING = _fake_streaming
    FAKE_STREAMING_INTERVAL = _fake_streaming_interval
    PASSWORD = _password
    MAX_REQUESTS_PER_MINUTE = _max_requests_per_minute
    MAX_REQUESTS_PER_DAY_PER_IP = _max_requests_per_day_per_ip

async def verify_user_agent(request: Request):
    if not settings.WHITELIST_USER_AGENT:
        return
    if request.headers.get("User-Agent") not in settings.WHITELIST_USER_AGENT:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Not allowed client")

# todo : 添加 gemini 支持(流式返回)
async def get_cache(cache_key,is_stream: bool,is_gemini=False):
    # 检查缓存是否存在，如果存在，返回缓存
    cached_response, cache_hit = await response_cache_manager.get_and_remove(cache_key)
    
    if cache_hit and cached_response:
        log('info', f"缓存命中: {cache_key[:8]}...", 
            extra={'request_type': 'non-stream', 'model': cached_response.model})
        
        if is_gemini:
            if is_stream:
                data = f"data: {json.dumps(cached_response.data, ensure_ascii=False)}\n\n"
                return StreamingResponse(data, media_type="text/event-stream")
            else:
                return cached_response.data
            
        
        if is_stream:
            chunk = openAI_from_Gemini(cached_response,stream=True)
            return StreamingResponse(chunk, media_type="text/event-stream")
        else: 
            return openAI_from_Gemini(cached_response,stream=False)

    return None

@router.get("/aistudio/models",response_model=ModelList)
async def aistudio_list_models(_ = Depends(custom_verify_password),
                               _2 = Depends(verify_user_agent)):
    if settings.WHITELIST_MODELS:
        filtered_models = [model for model in GeminiClient.AVAILABLE_MODELS if model in settings.WHITELIST_MODELS]
    else:
        filtered_models = [model for model in GeminiClient.AVAILABLE_MODELS if model not in settings.BLOCKED_MODELS]
    return ModelList(data=[{"id": model, "object": "model", "created": 1678888888, "owned_by": "organization-owner"} for model in filtered_models])

@router.get("/vertex/models",response_model=ModelList)
async def vertex_list_models(request: Request, 
                             _ = Depends(custom_verify_password),
                             _2 = Depends(verify_user_agent)):
    # 使用vertex/routes/models_api的实现
    return await models_api.list_models(request, current_api_key)

# API路由
@router.get("/v1/models",response_model=ModelList)
@router.get("/models",response_model=ModelList)
async def list_models(request: Request,
                      _ = Depends(custom_verify_password),
                      _2 = Depends(verify_user_agent)):
    if settings.ENABLE_VERTEX:
        return await vertex_list_models(request, _, _2)
    return await aistudio_list_models(_, _2)

@router.post("/aistudio/chat/completions", response_model=ChatCompletionResponse)
async def aistudio_chat_completions(
    request: Union[ChatCompletionRequest, AIRequest],
    http_request: Request,
    _ = Depends(custom_verify_password),
    _2 = Depends(verify_user_agent),
):
    format_type = getattr(request, 'format_type', None)
    if format_type and (format_type == "gemini"):
        is_gemini = True
    else:
        is_gemini = False
    
    # 生成缓存键 - 用于匹配请求内容对应缓存
    if settings.PRECISE_CACHE:
        cache_key = generate_cache_key(request, is_gemini = is_gemini)
    else:    
        cache_key = generate_cache_key(request, last_n_messages = settings.CALCULATE_CACHE_ENTRIES,is_gemini = is_gemini)
    
    # 请求前基本检查
    await protect_from_abuse(
        http_request, 
        settings.MAX_REQUESTS_PER_MINUTE, 
        settings.MAX_REQUESTS_PER_DAY_PER_IP)
    
    if request.model not in GeminiClient.AVAILABLE_MODELS:
        log('error', "无效的模型", 
            extra={'model': request.model, 'status_code': 400})
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="无效的模型")
    
    
    # 记录请求缓存键信息
    log('info', f"请求缓存键: {cache_key[:8]}...", 
        extra={'request_type': 'non-stream', 'model': request.model})
    
    # 检查缓存是否存在，如果存在，返回缓存
    cached_response = await get_cache(cache_key, is_stream = request.stream,is_gemini=is_gemini)
    if cached_response :
        return cached_response
    
    if not settings.PUBLIC_MODE:
        # 构建包含缓存键的活跃请求池键
        pool_key = f"{cache_key}"
        
        # 查找所有使用相同缓存键的活跃任务
        active_task = active_requests_manager.get(pool_key)
        if active_task and not active_task.done():
            log('info', f"发现相同请求的进行中任务", 
                extra={'request_type': 'stream' if request.stream else "non-stream", 'model': request.model})
            
            # 等待已有任务完成
            try:
                # 设置超时，避免无限等待
                await asyncio.wait_for(active_task, timeout=240)
                
                # 使用任务结果
                if active_task.done() and not active_task.cancelled():
                    
                    result = active_task.result()
                    active_requests_manager.remove(pool_key)
                    if result:
                        return result
            
            except (asyncio.TimeoutError, asyncio.CancelledError) as e:
                # 任务超时或被取消的情况下，记录日志然后让代码继续执行
                error_type = "超时" if isinstance(e, asyncio.TimeoutError) else "被取消"
                log('warning', f"等待已有任务{error_type}: {pool_key}", 
                    extra={'request_type': 'non-stream', 'model': request.model})
                
                # 从活跃请求池移除该任务
                if active_task.done() or active_task.cancelled():
                    active_requests_manager.remove(pool_key)
                    log('info', f"已从活跃请求池移除{error_type}任务: {pool_key}", 
                        extra={'request_type': 'non-stream'})
    
        
    if request.stream:
        # 流式请求处理任务
        process_task = asyncio.create_task(
            process_stream_request(
                chat_request = request, 
                key_manager=key_manager,
                response_cache_manager = response_cache_manager,
                safety_settings = safety_settings,
                safety_settings_g2 = safety_settings_g2,
                cache_key = cache_key
            )
        )
    
    else:
        # 检查是否启用非流式保活功能
        if settings.NONSTREAM_KEEPALIVE_ENABLED:
            # 使用带保活功能的非流式请求处理
            process_task = asyncio.create_task(
                process_nonstream_with_keepalive_stream(
                    chat_request = request,
                    key_manager = key_manager,
                    response_cache_manager = response_cache_manager,
                    safety_settings = safety_settings,
                    safety_settings_g2 = safety_settings_g2,
                    cache_key = cache_key,
                    is_gemini = is_gemini
                )
            )
        else:
            # 创建非流式请求处理任务
            process_task = asyncio.create_task(
                process_request(
                    chat_request = request,
                    key_manager = key_manager,
                    response_cache_manager = response_cache_manager,
                    safety_settings = safety_settings,
                    safety_settings_g2 = safety_settings_g2,
                    cache_key = cache_key
                )
            )

    if not settings.PUBLIC_MODE:
        # 将任务添加到活跃请求池
        active_requests_manager.add(pool_key, process_task)
    
    # 等待任务完成
    try:
        response = await process_task
        if not settings.PUBLIC_MODE:
            active_requests_manager.remove(pool_key)
        
        return response
    except Exception as e:
        if not settings.PUBLIC_MODE:
            # 如果任务失败，从活跃请求池中移除
            active_requests_manager.remove(pool_key)
        
        # 检查是否已有缓存的结果（可能是由另一个任务创建的）
        cached_response = await get_cache(cache_key, is_stream = request.stream,is_gemini=is_gemini)
        if cached_response :
            return cached_response
        
        # 发送错误信息给客户端
        raise HTTPException(status_code=500, detail=f" hajimi 服务器内部处理时发生错误\n具体原因:{e}")

@router.post("/vertex/chat/completions", response_model=ChatCompletionResponse)
async def vertex_chat_completions(
    request: ChatCompletionRequest, 
    http_request: Request,
    _dp = Depends(custom_verify_password),
    _du = Depends(verify_user_agent),
    ):
    # 使用vertex/routes/chat_api的实现
    
    # 转换消息格式
    openai_messages = []
    for message in request.messages:
        openai_messages.append(OpenAIMessage(
            role=message.get('role', ''),
            content=message.get('content', '')
        ))
    
    # 转换请求格式
    vertex_request = OpenAIRequest(
        model=request.model,
        messages=openai_messages,
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
        top_k=request.top_k,
        stream=request.stream,
        stop=request.stop,
        presence_penalty=request.presence_penalty,
        frequency_penalty=request.frequency_penalty,
        seed=getattr(request, 'seed', None),
        logprobs=getattr(request, 'logprobs', None),
        response_logprobs=getattr(request, 'response_logprobs', None),
        n=request.n
    )
    
    # 调用vertex/routes/chat_api的实现
    return await chat_api.chat_completions(http_request, vertex_request, current_api_key)

@router.post("/v1/chat/completions", response_model=ChatCompletionResponse)
@router.post("/chat/completions", response_model=ChatCompletionResponse)
async def chat_completions(
    request: ChatCompletionRequest,
    http_request: Request,
    _dp = Depends(custom_verify_password),
    _du = Depends(verify_user_agent),
):
    """处理API请求的主函数，根据需要处理流式或非流式请求"""
    if settings.ENABLE_VERTEX:
        return await vertex_chat_completions(request, http_request, _dp, _du)
    return await aistudio_chat_completions(request, http_request, _dp, _du)

@router.post("/gemini/{api_version:str}/models/{model_and_responseType:path}")
async def gemini_chat_completions(
    request: Request,
    model_and_responseType: str = Path(...),
    key: Optional[str] = Query(None),
    alt: Optional[str] = Query(None, description=" sse 或 None"),
    payload: ChatRequestGemini = Body(...),
    _dp = Depends(custom_verify_password),
    _du = Depends(verify_user_agent),
):
    # 提取路径参数
    is_stream = False
    try:
        model_name, action_type = model_and_responseType.split(":", 1)
        if action_type == "streamGenerateContent":
            is_stream = True
        
    except ValueError:
        raise HTTPException(status_code=400, detail="无效的请求路径")
    
    geminiRequest = AIRequest(payload=payload,model=model_name,stream=is_stream,format_type='gemini')
    return await aistudio_chat_completions(geminiRequest, request, _dp, _du)