xiaobai-review/xiaobai-datahub/datahub/siteconsole.py

from __future__ import annotations

from typing import Any

from datahub.serving import ApiError
from datahub.siteauth import SiteBridge, SiteBridgeError

VENDOR_PRESETS: tuple[dict[str, str], ...] = (
    {"id": "openai", "label": "OpenAI", "base_url": "https://api.openai.com/v1"},
    {"id": "deepseek", "label": "DeepSeek", "base_url": "https://api.deepseek.com/v1"},
    {"id": "moonshot", "label": "Moonshot", "base_url": "https://api.moonshot.cn/v1"},
    {"id": "dashscope", "label": "阿里云百炼", "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1"},
    {"id": "zhipu", "label": "智谱 GLM", "base_url": "https://open.bigmodel.cn/api/paas/v4"},
    {"id": "siliconflow", "label": "SiliconFlow", "base_url": "https://api.siliconflow.cn/v1"},
)


class SiteConsole:
    """Console-side view of the data the review site still owns.

    The model pool, member roster and invite codes live in the review site's
    database — this console reads and writes them over the bridge instead of
    copying them, so there is exactly one source of truth. Every method turns a
    bridge failure into an ``ApiError`` the console frontend already knows how
    to render.
    """

    def __init__(self, bridge: SiteBridge) -> None:
        self.bridge = bridge

    def _call(self, path: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        try:
            return self.bridge.call(path, payload)
        except SiteBridgeError as exc:
            # 主站因入参不合法而拒绝（Key 不对、模型 ID 不合法）是操作者的问题，
            # 照原样退回 400；只有主站真的不可达才算中枢侧不可用。
            code = "INVALID_ARGUMENT" if exc.caller_fault else "SOURCE_UNAVAILABLE"
            raise ApiError(code, str(exc)) from exc

    # ---------------------------------------------------------------- models
    def models(self) -> dict[str, Any]:
        payload = self._call("/api/hub-admin/status")
        llm = payload.get("llm") or {}
        models = list(llm.get("models") or [])
        return {
            "vendors": [dict(preset) for preset in VENDOR_PRESETS],
            "groups": _group_by_vendor(models),
            "primary_model_id": str(llm.get("primary_model_id") or ""),
            "fallback_model_id": str(llm.get("fallback_model_id") or ""),
            "models": models,
        }

    def save_models(self, body: dict[str, Any]) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        if "models" in body:
            payload["models"] = body.get("models") or []
        for key in ("primary_model_id", "fallback_model_id"):
            if key in body:
                payload[key] = str(body.get(key) or "")
        if not payload:
            raise ApiError("INVALID_ARGUMENT", "没有需要保存的模型配置")
        self._call("/api/hub-admin/settings/save", payload)
        return self.models()

    def test_model(self, body: dict[str, Any]) -> dict[str, Any]:
        payload = self._call(
            "/api/hub-admin/settings/test",
            {"model_id": str(body.get("model_id") or ""), "profile": body.get("profile") or {}},
        )
        return {"result": payload.get("result") or {}}

    def fetch_models(self, body: dict[str, Any]) -> dict[str, Any]:
        base_url = str(body.get("base_url") or "").strip()
        if not base_url:
            raise ApiError("INVALID_ARGUMENT", "请先填写供应商接口地址")
        payload = self._call(
            "/api/hub-admin/models/fetch",
            {"base_url": base_url, "api_key": str(body.get("api_key") or "")},
        )
        return {"models": payload.get("models") or []}

    # --------------------------------------------------------------- members
    def members(self) -> dict[str, Any]:
        payload = self._call("/api/hub-admin/members")
        return {
            "users": payload.get("users") or [],
            "membership": payload.get("membership") or {},
        }

    def save_member(self, body: dict[str, Any]) -> dict[str, Any]:
        payload = dict(body or {})
        if not payload:
            raise ApiError("INVALID_ARGUMENT", "没有需要保存的会员设置")
        self._call("/api/hub-admin/membership/save", payload)
        return self.members()

    def save_quota(self, body: dict[str, Any]) -> dict[str, Any]:
        """Daily call quota is a system setting, not a per-user membership row."""
        try:
            limit = int(body.get("member_daily_limit") or 0)
        except (TypeError, ValueError) as exc:
            raise ApiError("INVALID_ARGUMENT", "每日调用额度必须是整数") from exc
        if limit < 1:
            raise ApiError("INVALID_ARGUMENT", "每日调用额度至少为 1")
        self._call("/api/hub-admin/settings/save", {"member_daily_limit": limit})
        return self.members()

    # --------------------------------------------------------------- invites
    def invites(self) -> dict[str, Any]:
        payload = self._call("/api/hub-admin/invites")
        return {"summary": payload.get("summary") or {}, "codes": payload.get("codes") or []}

    def create_invites(self, body: dict[str, Any], created_by: int) -> dict[str, Any]:
        payload = self._call(
            "/api/hub-admin/invites/create",
            {
                "count": body.get("count") or 1,
                "note": str(body.get("note") or ""),
                "created_by": int(created_by or 0),
            },
        )
        # `created` carries the plaintext codes and is the only moment they are
        # ever returned; the list under `codes` is the masked roster.
        return {
            "created": payload.get("created") or [],
            "summary": payload.get("summary") or {},
            "codes": payload.get("codes") or [],
        }

    def revoke_invite(self, body: dict[str, Any]) -> dict[str, Any]:
        reference = str(body.get("code_id") or body.get("code") or "").strip()
        if not reference:
            raise ApiError("INVALID_ARGUMENT", "请选择要作废的邀请码")
        self._call("/api/hub-admin/invites/revoke", {"code_id": reference})
        return self.invites()


def _group_by_vendor(models: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Bucket a flat model list by base URL so the console can render vendors.

    The review site stores one row per model with its own base URL; the console
    shows vendors with their models nested, so the base URL is the grouping key
    and the preset table only supplies a friendly label when it recognises one.
    """
    labels = {preset["base_url"]: preset["label"] for preset in VENDOR_PRESETS}
    order: list[str] = []
    buckets: dict[str, list[dict[str, Any]]] = {}
    for model in models:
        base_url = str(model.get("base_url") or "").strip()
        if base_url not in buckets:
            buckets[base_url] = []
            order.append(base_url)
        buckets[base_url].append(model)
    groups: list[dict[str, Any]] = []
    for base_url in order:
        entries = buckets[base_url]
        configured = next((entry for entry in entries if entry.get("api_key_last4")), None)
        groups.append(
            {
                "base_url": base_url,
                "label": labels.get(base_url) or _vendor_label(base_url),
                "configured": any(entry.get("configured") for entry in entries),
                "key_last4": str((configured or {}).get("api_key_last4") or ""),
                "models": entries,
            }
        )
    return groups


def _vendor_label(base_url: str) -> str:
    host = base_url.split("//")[-1].split("/")[0]
    return host or "自定义供应商"