Aether/_deprecated_py_src/services/model/upstream_fetcher.py

"""
上游模型获取公共模块

提供从上游 API 获取模型列表的公共函数；通用 api_format 抓取能力统一来自 core.api_format.capabilities，供以下场景使用：
- 定时任务自动获取（fetch_scheduler.py）
- 管理后台手动查询（provider_query.py）
"""

import asyncio
import json
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any, Awaitable, Callable

import httpx

from src.config.settings import config
from src.core.api_format import get_extra_headers_from_endpoint
from src.core.api_format.capabilities import BROWSER_FINGERPRINT_HEADERS
from src.core.api_format.headers import build_adapter_headers_for_endpoint
from src.core.logger import logger

# 并发请求限制
MAX_CONCURRENT_REQUESTS = 5

# 模型获取格式优先级：同族内优先使用 chat 端点，若无则回退到 cli 端点
MODEL_FETCH_FORMAT_PRIORITY: list[tuple[str, ...]] = [
    ("openai:chat", "openai:cli", "openai:compact"),
    ("claude:chat", "claude:cli"),
    ("gemini:chat", "gemini:cli"),
]

# Return tuple signature:
# (models, errors, has_success, upstream_metadata)
_ModelsFetcher = Callable[
    ["UpstreamModelsFetchContext", float],
    Awaitable[tuple[list[dict], list[str], bool, dict[str, Any] | None]],
]


@dataclass(frozen=True, slots=True)
class EndpointFetchConfig:
    """端点获取配置（纯数据，不依赖 DB session）。

    从 ProviderEndpoint ORM 对象提取必要字段，确保在 DB session 关闭后
    仍可安全使用（避免 DetachedInstanceError）。
    """

    base_url: str
    extra_headers: dict[str, str] | None = None


def build_format_to_config(endpoints: Iterable[Any]) -> dict[str, EndpointFetchConfig]:
    """将活跃的 ProviderEndpoint 转换为 api_format -> EndpointFetchConfig 映射。

    应在 DB session 活跃时调用，提取 ORM 对象上的 base_url 和 header_rules，
    转换为 session 无关的纯数据结构。
    """
    result: dict[str, EndpointFetchConfig] = {}
    for ep in endpoints:
        if not getattr(ep, "is_active", False):
            continue
        result[ep.api_format] = EndpointFetchConfig(
            base_url=ep.base_url,
            extra_headers=get_extra_headers_from_endpoint(ep),
        )
    return result


@dataclass(frozen=True)
class UpstreamModelsFetchContext:
    """上游模型获取上下文（Key 级别）。"""

    provider_type: str
    api_key_value: str
    format_to_endpoint: dict[str, EndpointFetchConfig]
    proxy_config: dict[str, Any] | None = None
    auth_config: dict[str, Any] | None = None


class UpstreamModelsFetcherRegistry:
    """按 provider_type 注册上游模型获取策略，避免到处写特判。"""

    _fetchers: dict[str, _ModelsFetcher] = {}

    @classmethod
    def register(cls, *, provider_types: list[str], fetcher: _ModelsFetcher) -> None:
        for pt in provider_types:
            if not pt:
                continue
            cls._fetchers[pt.lower()] = fetcher

    @classmethod
    def get(cls, provider_type: str) -> _ModelsFetcher | None:
        if not provider_type:
            return None
        return cls._fetchers.get(provider_type.lower())


async def _fetch_models_default(
    ctx: UpstreamModelsFetchContext,
    timeout_seconds: float,
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
    endpoint_configs = build_all_format_configs(ctx.api_key_value, ctx.format_to_endpoint)
    models, errors, has_success = await fetch_models_from_endpoints(
        endpoint_configs, timeout=timeout_seconds, proxy_config=ctx.proxy_config
    )
    return models, errors, has_success, None


async def fetch_models_for_key(
    ctx: UpstreamModelsFetchContext,
    *,
    timeout_seconds: float = 30.0,
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
    """统一入口：按 provider_type 选择策略获取模型列表（可附带 upstream_metadata）。"""
    # Ensure provider plugins (including custom model fetchers) are registered.
    from src.services.provider.envelope import ensure_providers_bootstrapped

    ensure_providers_bootstrapped(provider_types=[ctx.provider_type] if ctx.provider_type else None)

    fetcher = UpstreamModelsFetcherRegistry.get(ctx.provider_type) or _fetch_models_default
    return await fetcher(ctx, timeout_seconds)


def merge_upstream_metadata(
    current: dict[str, Any] | None,
    incoming: dict[str, Any],
) -> dict[str, Any]:
    """合并上游元数据，对 quota_by_model 做模型级深度合并。

    以上游返回为准（全量替换），不再保留上游已下架的模型。
    仅对上游返回的模型补充旧数据中的 reset_time（当新数据缺少时）。
    """
    merged: dict[str, Any] = dict(current) if isinstance(current, dict) else {}
    for ns_key, ns_val in incoming.items():
        old_ns = merged.get(ns_key)
        if (
            isinstance(ns_val, dict)
            and isinstance(old_ns, dict)
            and "quota_by_model" in ns_val
            and "quota_by_model" in old_ns
        ):
            old_qbm = old_ns["quota_by_model"]
            new_qbm = ns_val["quota_by_model"]
            if isinstance(old_qbm, dict) and isinstance(new_qbm, dict):
                # 保留新数据中已有模型的旧 reset_time
                for model_id, new_info in new_qbm.items():
                    if not isinstance(new_info, dict):
                        continue
                    old_info = old_qbm.get(model_id)
                    if (
                        isinstance(old_info, dict)
                        and "reset_time" in old_info
                        and "reset_time" not in new_info
                    ):
                        new_info["reset_time"] = old_info["reset_time"]
        merged[ns_key] = ns_val
    return merged


def build_all_format_configs(
    api_key_value: str,
    format_to_endpoint: dict[str, EndpointFetchConfig],
) -> list[dict]:
    """
    构建所有 API 格式的端点配置

    只对实际配置了端点的格式构建请求配置，不同端点的 base_url 可能不同，
    不应使用某个端点的 base_url 去尝试其他格式。

    Args:
        api_key_value: 解密后的 API Key
        format_to_endpoint: API 格式到 EndpointFetchConfig 的映射

    Returns:
        端点配置列表，每个配置包含 api_key, base_url, api_format, extra_headers
    """
    if not format_to_endpoint:
        return []

    # 同族内优先使用 chat 端点，若无则回退到 cli 端点
    configs: list[dict] = []
    for candidates in MODEL_FETCH_FORMAT_PRIORITY:
        fmt = next((f for f in candidates if f in format_to_endpoint), None)
        if fmt is not None:
            cfg = format_to_endpoint[fmt]

            base_url = str(getattr(cfg, "base_url", "") or "")
            extra_headers: dict[str, str] | None

            if isinstance(cfg, EndpointFetchConfig):
                extra_headers = cfg.extra_headers
            else:
                # 允许直接传递类似 ProviderEndpoint 的对象（测试/独立使用场景）。
                candidate_extra = getattr(cfg, "extra_headers", None)
                if isinstance(candidate_extra, dict):
                    extra_headers = {str(k): str(v) for k, v in candidate_extra.items() if k}
                else:
                    extra_headers = get_extra_headers_from_endpoint(cfg)

            configs.append(
                {
                    "api_key": api_key_value,
                    "base_url": base_url,
                    "api_format": fmt,
                    "extra_headers": extra_headers,
                }
            )
    return configs


async def _build_models_proxy_snapshot(
    proxy_config: dict[str, Any] | None,
) -> Any:
    from src.services.request.execution_runtime_plan import build_proxy_snapshot

    return await build_proxy_snapshot(proxy_config, label="upstream model fetch")


def _build_model_fetch_url(
    api_format: str,
    base_url: str,
    api_key: str,
    after_id: str | None = None,
) -> str:
    normalized_api_format = str(api_format or "").strip().lower()
    base_url = str(base_url or "").rstrip("/")

    if normalized_api_format.startswith("gemini:"):
        if base_url.endswith("/v1beta"):
            return f"{base_url}/models?key={api_key}"
        return f"{base_url}/v1beta/models?key={api_key}"

    if base_url.endswith("/v1"):
        url = f"{base_url}/models"
    else:
        url = f"{base_url}/v1/models"

    if normalized_api_format.startswith("claude:"):
        if after_id:
            return f"{url}?limit=100&after_id={after_id}"
        return f"{url}?limit=100"

    return url


def _build_model_fetch_headers(
    api_format: str,
    api_key: str,
    extra_headers: dict[str, str] | None,
) -> dict[str, str]:
    normalized_api_format = str(api_format or "").strip().lower()

    if normalized_api_format == "openai:cli":
        headers = {"User-Agent": config.internal_user_agent_openai_cli}
        if extra_headers:
            headers.update(extra_headers)
        return build_adapter_headers_for_endpoint("openai:chat", api_key, headers)

    if normalized_api_format == "openai:compact":
        headers = {"User-Agent": config.internal_user_agent_openai_cli}
        if extra_headers:
            headers.update(extra_headers)
        return build_adapter_headers_for_endpoint("openai:chat", api_key, headers)

    if normalized_api_format == "claude:cli":
        headers = {"User-Agent": config.internal_user_agent_claude_cli}
        if extra_headers:
            headers.update(extra_headers)
        return build_adapter_headers_for_endpoint(normalized_api_format, api_key, headers)

    if normalized_api_format == "claude:chat":
        headers = build_adapter_headers_for_endpoint(normalized_api_format, api_key, extra_headers)
        if "authorization" not in {str(k).lower() for k in headers}:
            headers["Authorization"] = f"Bearer {api_key}"
        return headers

    if normalized_api_format == "gemini:cli":
        headers = {
            **BROWSER_FINGERPRINT_HEADERS,
            "User-Agent": config.internal_user_agent_gemini_cli,
        }
        if extra_headers:
            headers.update(extra_headers)
        return headers

    if normalized_api_format == "gemini:chat":
        headers = {**BROWSER_FINGERPRINT_HEADERS}
        if extra_headers:
            headers.update(extra_headers)
        return headers

    return build_adapter_headers_for_endpoint(normalized_api_format, api_key, extra_headers)


def _parse_model_fetch_response(
    api_format: str,
    payload: Any,
) -> tuple[list[dict[str, Any]], bool, str | None]:
    normalized_api_format = str(api_format or "").strip().lower()

    if normalized_api_format.startswith("gemini:"):
        if isinstance(payload, dict) and isinstance(payload.get("models"), list):
            out: list[dict[str, Any]] = []
            for model in payload["models"]:
                if not isinstance(model, dict):
                    continue
                out.append(
                    {
                        "id": str(model.get("name", "")).replace("models/", ""),
                        "owned_by": "google",
                        "display_name": model.get("displayName", ""),
                        "api_format": normalized_api_format,
                    }
                )
            return out, False, None
        return [], False, None

    page_models: list[dict[str, Any]] = []
    has_more = False
    if isinstance(payload, dict) and isinstance(payload.get("data"), list):
        page_models = [m for m in payload["data"] if isinstance(m, dict)]
        has_more = bool(payload.get("has_more"))
    elif isinstance(payload, list):
        page_models = [m for m in payload if isinstance(m, dict)]

    for model in page_models:
        model.setdefault("api_format", normalized_api_format)

    next_cursor = None
    if normalized_api_format.startswith("claude:") and isinstance(payload, dict):
        raw_last_id = payload.get("last_id")
        if has_more and isinstance(raw_last_id, str) and raw_last_id.strip():
            next_cursor = raw_last_id.strip()

    return page_models, has_more, next_cursor


def _extract_rust_error_message(result: Any) -> str:
    payload = getattr(result, "response_json", None)
    if isinstance(payload, dict):
        err = payload.get("error")
        if isinstance(err, dict):
            message = str(err.get("message") or "").strip()
            if message:
                return message[:500]
        message = str(payload.get("message") or "").strip()
        if message:
            return message[:500]
        try:
            return json.dumps(payload, ensure_ascii=False)[:500]
        except Exception:
            pass

    body_bytes = getattr(result, "response_body_bytes", None)
    if body_bytes:
        try:
            return body_bytes.decode("utf-8", errors="replace")[:500]
        except Exception:
            return "(binary body)"
    return "(empty)"


async def _try_rust_fetch_models_for_api_format(
    *,
    api_format: str,
    base_url: str,
    api_key: str,
    extra_headers: dict[str, str] | None,
    timeout: float,
    proxy_config: dict[str, Any] | None,
) -> tuple[list[dict[str, Any]], str | None, bool] | None:
    from src.services.request.execution_runtime_plan import (
        ExecutionPlan,
        ExecutionPlanBody,
        ExecutionPlanTimeouts,
    )
    from src.services.request.execution_runtime_client import (
        ExecutionRuntimeClient,
        ExecutionRuntimeClientError,
    )

    if config.execution_runtime_backend != "rust":
        return None

    try:
        proxy_snapshot = await _build_models_proxy_snapshot(proxy_config)
        headers = _build_model_fetch_headers(api_format, api_key, extra_headers)
        models: list[dict[str, Any]] = []
        next_cursor: str | None = None
        seen_ids: set[str] = set()

        for _ in range(20):
            url = _build_model_fetch_url(api_format, base_url, api_key, next_cursor)
            result = await ExecutionRuntimeClient().execute_sync_json(
                ExecutionPlan(
                    request_id=f"model-fetch:{api_format}:{next_cursor or 'root'}",
                    candidate_id=None,
                    provider_name=str(api_format.split(":", 1)[0] or "unknown"),
                    provider_id="",
                    endpoint_id="",
                    key_id="",
                    method="GET",
                    url=url,
                    headers=dict(headers),
                    body=ExecutionPlanBody(),
                    stream=False,
                    provider_api_format=api_format,
                    client_api_format=api_format,
                    model_name="models",
                    proxy=proxy_snapshot,
                    timeouts=ExecutionPlanTimeouts(
                        connect_ms=min(int(timeout * 1000), 30_000),
                        read_ms=int(timeout * 1000),
                        write_ms=int(timeout * 1000),
                        pool_ms=min(int(timeout * 1000), 30_000),
                        total_ms=int(timeout * 1000),
                    ),
                )
            )

            if result.status_code != 200:
                error_body = _extract_rust_error_message(result)
                return [], f"HTTP {result.status_code}: {error_body}", False

            payload = result.response_json
            page_models, has_more, next_after_id = _parse_model_fetch_response(api_format, payload)
            for model in page_models:
                model_id = model.get("id")
                if isinstance(model_id, str) and model_id and model_id in seen_ids:
                    continue
                if isinstance(model_id, str) and model_id:
                    seen_ids.add(model_id)
                models.append(model)

            if not has_more or not next_after_id or next_after_id == next_cursor:
                return models, None, True
            next_cursor = next_after_id

        return models, None, True
    except (ExecutionRuntimeClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
        logger.warning("Rust model fetch fallback for {}: {}", api_format, exc)
        return None
    except Exception as exc:
        logger.warning("Rust model fetch unexpected fallback for {}: {}", api_format, exc)
        return None


async def fetch_models_from_endpoints(
    endpoint_configs: list[dict],
    timeout: float = 30.0,
    proxy_config: dict[str, Any] | None = None,
) -> tuple[list[dict], list[str], bool]:
    """
    从多个端点并发获取模型

    Args:
        endpoint_configs: 端点配置列表，每个配置包含 api_key, base_url, api_format, extra_headers
        timeout: 请求超时时间（秒）
        proxy_config: 代理配置（可选），支持系统默认回退

    Returns:
        (模型列表, 错误列表, 是否有成功)
    """
    all_models: list[dict] = []
    errors: list[str] = []
    has_success = False
    semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)

    async def fetch_one(config: dict) -> tuple[list, str | None, bool]:
        base_url = config["base_url"]
        if not base_url:
            return [], None, False
        base_url = base_url.rstrip("/")
        api_format = config["api_format"]
        api_key_value = config["api_key"]
        extra_headers = config.get("extra_headers")

        try:
            async with semaphore:
                rust_result = await _try_rust_fetch_models_for_api_format(
                    api_format=api_format,
                    base_url=base_url,
                    api_key=api_key_value,
                    extra_headers=extra_headers,
                    timeout=timeout,
                    proxy_config=proxy_config,
                )
                if rust_result is None:
                    return [], f"{api_format}: rust executor unavailable", False
                models, error, success = rust_result

            for m in models:
                if "api_format" not in m:
                    m["api_format"] = api_format

            # 即使返回空列表，只要没有错误也算成功
            return models, error, success
        except httpx.TimeoutException:
            logger.warning("获取 {} 模型超时", api_format)
            return [], f"{api_format}: timeout", False
        except Exception:
            logger.exception("获取 {} 模型出错", api_format)
            return [], f"{api_format}: error", False

    results = await asyncio.gather(*[fetch_one(c) for c in endpoint_configs])
    for models, error, success in results:
        all_models.extend(models)
        if error:
            errors.append(error)
        if success:
            has_success = True

    return all_models, errors, has_success