Aether/_deprecated_py_src/api/handlers/gemini/handler.py

"""
Gemini Chat Handler

处理 Gemini API 格式的请求
"""

from __future__ import annotations

from typing import Any

from starlette.requests import Request

from src.api.handlers.base.chat_handler_base import ChatHandlerBase
from src.core.api_format import ApiFamily, EndpointKind


class GeminiChatHandler(ChatHandlerBase):
    """
    Gemini Chat Handler - 处理 Google Gemini API 格式的请求

    格式特点:
    - 使用 promptTokenCount / candidatesTokenCount
    - 支持 cachedContentTokenCount
    - 请求格式: GeminiRequest
    - 响应格式: JSON 数组流（非 SSE）
    """

    FORMAT_ID = "gemini:chat"
    API_FAMILY = ApiFamily.GEMINI
    ENDPOINT_KIND = EndpointKind.CHAT

    async def _resolve_preferred_key_ids(
        self,
        model_name: str,  # noqa: ARG002 - 仅做文件绑定
        request_body: dict[str, Any] | None = None,
    ) -> list[str] | None:
        """
        从 files/xxx 绑定关系中解析优先 Key ID 列表。

        Gemini 文件与上传它的 API Key 绑定，必须使用同一 Key 访问。
        此方法从缓存中查找文件→Key 映射，优先使用正确的 Key。

        当同一源文件被上传到多个 Key 时，会返回所有可用的 Key ID，
        让系统能够选择任意可用的 Key。

        注意事项：
        - 如果映射缺失（缓存过期/重启），会记录警告，请求可能失败
        - 优先返回所有支持该文件的 Key，让调度器选择可用的
        """
        from src.core.logger import logger
        from src.services.gemini_files_mapping import (
            extract_file_names_from_request,
            get_all_key_ids_for_file,
        )

        file_names = extract_file_names_from_request(request_body or {})
        if not file_names:
            return None

        all_key_ids: set[str] = set()
        unmapped_files: list[str] = []

        for file_name in file_names:
            # 获取所有支持该文件的 Key（包括通过 source_hash 关联的）
            key_ids = await get_all_key_ids_for_file(file_name)
            if key_ids:
                all_key_ids.update(key_ids)
            else:
                unmapped_files.append(file_name)

        # 警告：映射缺失
        if unmapped_files:
            logger.warning(
                f"[{self.request_id}] Gemini 文件→Key 映射缺失: {unmapped_files}，"
                "请求可能失败（文件属于其他 Key 或映射已过期）"
            )

        if all_key_ids:
            logger.debug(f"[{self.request_id}] 文件引用可用的 Key: {list(all_key_ids)}")

        return list(all_key_ids) if all_key_ids else None

    def extract_model_from_request(
        self,
        request_body: dict[str, Any],
        path_params: dict[str, Any] | None = None,
    ) -> str:
        """
        从请求中提取模型名 - Gemini Chat 格式实现

        Gemini Chat 模式下，model 在请求体中（经过转换后的 GeminiRequest）。
        与 Gemini CLI 不同，CLI 模式的 model 在 URL 路径中。

        Args:
            request_body: 请求体
            path_params: URL 路径参数（Chat 模式通常不使用）

        Returns:
            模型名
        """
        # 优先从请求体获取，其次从 path_params
        model = request_body.get("model")
        if model:
            return str(model)
        if path_params and "model" in path_params:
            return str(path_params["model"])
        return "unknown"

    async def _convert_request(self, request: Request) -> None:
        """
        将请求转换为 Gemini 格式的 Pydantic 对象

        注意：此方法只做类型转换（dict → Pydantic），不做跨格式转换。
        跨格式转换由调度/执行层（TaskService + RequestDispatcher）在选中候选后、发送请求前执行，
        并受全局开关和端点配置控制。

        Args:
            request: 原始请求对象（应已是 Gemini 格式）

        Returns:
            GeminiRequest 对象
        """
        from src.models.gemini import GeminiRequest

        # 如果已经是 Gemini 格式 Pydantic 对象，直接返回
        if isinstance(request, GeminiRequest):
            return request

        # 如果是字典，转换为 Pydantic 对象（假设已是 Gemini 格式）
        if isinstance(request, dict):
            return GeminiRequest(**request)

        return request

    def _extract_usage(self, response: dict) -> dict[str, int]:
        """
        从 Gemini 响应中提取 token 使用情况

        调用 GeminiStreamParser.extract_usage 作为单一实现源
        """
        from src.api.handlers.gemini.stream_parser import GeminiStreamParser

        usage = GeminiStreamParser().extract_usage(response)

        if not usage:
            return {
                "input_tokens": 0,
                "output_tokens": 0,
                "cache_creation_input_tokens": 0,
                "cache_read_input_tokens": 0,
            }

        return {
            "input_tokens": usage.get("input_tokens", 0),
            "output_tokens": usage.get("output_tokens", 0),
            "cache_creation_input_tokens": 0,  # Gemini 不区分缓存创建
            "cache_read_input_tokens": usage.get("cached_tokens", 0),
        }

    def finalize_provider_request(
        self,
        request_body: dict[str, Any],
        *,
        mapped_model: str | None,
        provider_api_format: str | None,  # noqa: ARG002
    ) -> dict[str, Any]:
        from src.api.handlers.gemini.image_gen import (
            adapt_request_for_image_gen,
            is_image_gen_model,
        )

        # Sanitize Gemini contents: strip parts without a valid data-oneof
        # field and merge consecutive same-role entries.  This catches cases
        # missed by the normalizer (passthrough) or the antigravity envelope.
        from src.core.api_format.conversion.normalizers.gemini import (
            compact_gemini_contents,
        )

        contents = request_body.get("contents")
        if isinstance(contents, list):
            request_body["contents"] = compact_gemini_contents(contents)

        if not is_image_gen_model(mapped_model):
            return request_body
        return adapt_request_for_image_gen(request_body)

    def _normalize_response(self, response: dict) -> dict:
        """
        规范化 Gemini 响应

        Args:
            response: 原始响应

        Returns:
            规范化后的响应
        """
        # 作为中转站，直接透传响应，不做标准化处理
        return response