app/api/v2_endpoints.py

# -*- coding: utf-8 -*-
"""
处理 Gemini 原生 API (v2) 相关的路由。
"""
import logging # 导入日志模块
from fastapi import APIRouter, Request, HTTPException, status, Depends, Path # 导入 FastAPI 相关组件
from fastapi.responses import JSONResponse # 导入 JSONResponse
from typing import Dict, Any, Optional # 导入类型提示

# 导入自定义模块
from app.core.services.gemini import GeminiClient # 导入 GeminiClient (新路径)
from app.core.utils.response_wrapper import ResponseWrapper # 导入响应包装类 (新路径)
from app.core.context.store import load_context, save_context, convert_openai_to_gemini_contents, convert_gemini_to_storage_format, load_context_as_gemini # 导入上下文管理函数和转换函数 (新路径)
from app.core.keys.manager import APIKeyManager # 导入类型 (新路径)
import httpx # 导入 httpx 用于类型提示
# 导入依赖注入函数
from app.core.dependencies import get_key_manager, get_http_client
# 导入请求工具函数和处理工具函数
from app.core.utils.request_helpers import get_client_ip, get_current_timestamps # (新路径)
from app.core.processing.utils import check_rate_limits_and_update_counts, update_token_counts # (新路径)
from app import config as app_config # 导入应用配置
from app.core.database.utils import IS_MEMORY_DB # 导入数据库模式标志 (新路径)
from app.config import ENABLE_CONTEXT_COMPLETION # 导入全局上下文补全配置

# 导入认证依赖项
from app.api.middleware import verify_proxy_key # 导入代理 Key 验证依赖

# 导入 Pydantic 模型
from app.api.models import GeminiGenerateContentRequestV2, GeminiGenerateContentResponseV2 # 导入 v2 请求和响应模型

logger = logging.getLogger('my_logger') # 获取日志记录器实例

v2_router = APIRouter() # 创建 APIRouter 实例

async def _load_and_inject_context_v2(
    proxy_key: str,
    request_body: GeminiGenerateContentRequestV2,
    enable_context: bool
):
    """
    加载历史上下文并注入到请求体中。
    """
    if not enable_context:
        logger.debug(f"Key {proxy_key[:8]}... 的上下文补全已禁用，跳过加载和注入。")
        return

    # 加载历史上下文并转换为 Gemini 格式
    gemini_context = await load_context_as_gemini(proxy_key) # 加载并转换历史上下文
    if gemini_context: # 如果存在上下文历史
        # 将上下文注入到当前请求内容的开头
        request_body.contents = gemini_context + request_body.contents # 注入上下文
        logger.debug(f"为 Key {proxy_key[:8]}... 注入了 {len(gemini_context)} 条上下文消息。") # 记录注入的上下文数量
    else:
        logger.debug(f"Key {proxy_key[:8]}... 没有找到上下文或加载失败，跳过注入。") # 没有找到上下文或加载失败


async def _save_context_after_v2_success(
    proxy_key: str,
    original_contents: list,
    gemini_response: Dict[str, Any],
    enable_context: bool
):
    """
    在 v2 API 调用成功后保存上下文。
    """
    if not enable_context or not gemini_response or not gemini_response.get('candidates'):
        if enable_context:
             logger.debug(f"为 Key {proxy_key[:8]}... 跳过上下文存储：上下文补全未启用或响应无效。")
        return

    # 提取用户请求内容和模型响应内容
    # 假设原始请求内容的最后一条是用户消息
    user_message_gemini = original_contents[-1] if original_contents else None # 获取用户消息
    # 假设响应的第一个候选的 content 是模型响应
    model_response_gemini = gemini_response['candidates'][0]['content'] if gemini_response.get('candidates') else None # 获取模型响应

    if user_message_gemini and model_response_gemini: # 如果用户消息和模型响应都存在
        # 将 Gemini 格式的用户请求和模型响应转换为存储格式 (OpenAI 格式)
        new_context_entry = convert_gemini_to_storage_format(user_message_gemini, model_response_gemini) # 转换格式
        if new_context_entry: # 如果转换成功
            # 加载现有上下文，添加新条目，然后保存
            existing_context = await load_context(proxy_key) # 重新加载现有上下文
            updated_context = (existing_context or []) + new_context_entry # 合并新旧上下文
            await save_context(proxy_key, updated_context) # 保存更新后的上下文
            logger.debug(f"为 Key {proxy_key[:8]}... 存储了新的上下文回合。") # 存储了新的上下文回合
        else:
            logger.warning(f"为 Key {proxy_key[:8]}... 转换 Gemini 请求/响应为存储格式失败，跳过上下文存储。") # 转换 Gemini 请求/响应为存储格式失败
    else:
         logger.warning(f"为 Key {proxy_key[:8]}... 存储上下文失败：无法提取用户消息或模型响应。") # 存储上下文失败：无法提取用户消息或模型响应


@v2_router.post("/models/{model}:generateContent", response_model=GeminiGenerateContentResponseV2) # 定义 POST 请求端点，指定响应模型
async def generate_content_v2(
    request: Request, # FastAPI 请求对象
    request_body: GeminiGenerateContentRequestV2, # 请求体，使用 Pydantic 模型进行验证
    model: str = Path(..., description="要使用的 Gemini 模型名称"), # 从路径参数获取模型名称
    auth_data: Dict[str, Any] = Depends(verify_proxy_key), # 依赖项：验证代理 Key 并获取 Key 和配置
    key_manager: APIKeyManager = Depends(get_key_manager), # 注入 Key Manager
    http_client: httpx.AsyncClient = Depends(get_http_client), # 注入 HTTP Client
):
    """
    处理 Gemini 原生 API 的 generateContent 请求 (/v2)。
    """
    proxy_key = auth_data.get("key") # 从认证数据中获取代理 Key
    key_config = auth_data.get("config", {}) # 从认证数据中获取 Key 配置，默认为空字典
    enable_context = key_config.get('enable_context_completion', ENABLE_CONTEXT_COMPLETION) # 获取 Key 的上下文补全配置，如果 Key 配置中没有则使用全局配置

    logger.info(f"收到 /v2/models/{model}:generateContent 请求，使用 Key: {proxy_key[:8]}..., 上下文补全: {enable_context}") # 记录收到的请求信息

    # 初始化 Gemini 客户端，传入共享的 http_client
    client = GeminiClient(api_key=proxy_key, http_client=http_client) # 使用代理 Key 初始化 Gemini 客户端

    # 获取并注入上下文 (如果启用) - 委托给辅助函数
    original_contents = request_body.contents # 保存原始请求内容
    await _load_and_inject_context_v2(proxy_key, request_body, enable_context)


    # --- 获取客户端 IP 和时间戳 ---
    client_ip = get_client_ip(request) # 获取客户端 IP
    _, today_date_str_pt = get_current_timestamps() # 获取 PT 日期字符串

    # --- 速率限制检查 ---
    limits = app_config.MODEL_LIMITS.get(model) # 获取模型限制
    if not check_rate_limits_and_update_counts(proxy_key, model, limits): # 检查并更新速率限制计数
        # 如果达到限制，check_rate_limits_and_update_counts 会记录警告
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS, # 429 状态码
            detail=f"API Key for model '{model}' has reached rate limits. Please try again later." # 错误详情
        )

    # 调用 Gemini API
    try:
        # 将 Pydantic 模型转换为字典，以便传递给 GeminiClient
        request_payload = request_body.model_dump(exclude_none=True) # 转换为字典，排除 None 值
        logger.debug(f"调用 Gemini API，模型: {model}, Payload: {request_payload}") # 调用 Gemini API
        # 调用 GeminiClient 的 generate_content 方法
        gemini_response = await client.generate_content(model_name=model, request_payload=request_payload) # 调用 Gemini API
        logger.debug(f"收到 Gemini API 响应: {gemini_response}") # 收到 Gemini API 响应

        # --- 更新 Token 计数 ---
        prompt_tokens = None # 初始化 prompt_tokens
        if isinstance(gemini_response, dict) and 'usageMetadata' in gemini_response: # 检查响应是否为字典且包含 usageMetadata
            prompt_tokens = gemini_response['usageMetadata'].get('promptTokenCount') # 获取 promptTokenCount
        update_token_counts(proxy_key, model, limits, prompt_tokens, client_ip, today_date_str_pt) # 更新 token 计数

        # 存储上下文 (如果启用) - 委托给辅助函数
        await _save_context_after_v2_success(proxy_key, original_contents, gemini_response, enable_context)

        # 返回 Gemini 响应
        # 直接返回原始 Gemini 响应，不进行 OpenAI 包装
        return JSONResponse(content=gemini_response) # 返回 JSON 响应

    except HTTPException as e:
        # 捕获 FastAPI HTTPException 并重新抛出
        raise e
    except Exception as e:
        # 捕获其他所有异常，记录并返回 500 错误
        logger.error(f"处理 /v2/models/{model}:generateContent 请求时发生错误 (Key: {proxy_key[:8]}...): {e}", exc_info=True) # 处理 /v2/models/{model}:generateContent 请求时发生错误
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, # 内部服务器错误状态码
            detail=f"处理请求时发生内部错误: {e}" # 错误详情
        )