task_worker_v2_portable/config.py

import json
import os
import sys
from typing import Optional

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_CONFIG_NAME = "worker_config.json"

PLATFORM_DEFS = {
    "ds": {"name": "DeepSeek", "intercept_url": "chat/completion", "intercept_domain": "deepseek.com"},
    "kimi": {"name": "Kimi", "intercept_url": "ChatService/Chat", "intercept_domain": "kimi.com"},
    "qianwen": {"name": "通义千问", "intercept_url": "api/v2/chat", "intercept_domain": "qianwen.com"},
    "wenxin": {"name": "文心一言", "intercept_url": "aichat/api/conversation", "intercept_domain": "chat.baidu.com"},
    "doubao": {"name": "豆包", "intercept_url": "chat/completion", "intercept_domain": "doubao.com"},
}

API_TO_INTERNAL = {
    "deepseek": "ds", "kimi": "kimi", "tongyi": "qianwen",
    "doubao": "doubao", "baiduai": "wenxin",
    "yuanbao": None, "antafu": None,
}
INTERNAL_TO_API = {v: k for k, v in API_TO_INTERNAL.items() if v}
SUPPORTED_API_PLATFORMS = [k for k, v in API_TO_INTERNAL.items() if v]

COOKIE_PLATFORM_IDS = {"ds": "1", "qianwen": "2", "kimi": "4", "doubao": "5", "wenxin": "6"}

COOKIE_DOMAINS = {
    "ds": ".deepseek.com", "qianwen": ".qianwen.com", "kimi": ".kimi.com",
    "doubao": ".doubao.com", "wenxin": ".baidu.com",
}

QUOTA_KEYWORDS = (
    "免费次数用完", "次数已用完", "暂时无法使用", "开通.*专业版",
    "太多了", "有点累了", "算力不足", "耐心等待", "前往升级",
    "请求过于频繁", "稍后再试", "服务繁忙",
)


def _parse_bool(v) -> bool:
    if isinstance(v, bool):
        return v
    if isinstance(v, (int, float)):
        return bool(v)
    s = str(v).strip().lower()
    if s in ("1", "true", "yes", "on", "enable", "enabled", "开", "启用"):
        return True
    if s in ("0", "false", "no", "off", "disable", "disabled", "关", "停用", "停止"):
        return False
    raise ValueError(f"无法解析布尔值: {v!r}")


def _default_switches():
    return {p: True for p in SUPPORTED_API_PLATFORMS}


def normalize_platform_name(name: str) -> Optional[str]:
    n = (name or "").strip().lower()
    aliases = {
        "ds": "deepseek", "deepseek": "deepseek", "kimi": "kimi",
        "tongyi": "tongyi", "qianwen": "tongyi", "qw": "tongyi",
        "doubao": "doubao", "db": "doubao",
        "baiduai": "baiduai", "wenxin": "baiduai", "wx": "baiduai",
        "yuanbao": "yuanbao", "antafu": "antafu",
    }
    return aliases.get(n, n if n in API_TO_INTERNAL else None)


def default_config() -> dict:
    return {
        "api_base": "https://api.granking.com",
        "spider_token": "",
        "vendor_token": "",
        "cookie_pool": {
            "session_api": "http://granking-api.neicela.com/api/third",
            "app_id": "aa65700299848d6f21b969dbc9f6cf7c",
            "secret": "5588071d36f0bc61af849c311a03f2c4",
            "auth_token": "",
            "retries": 3,
        },
        "platform_switches": _default_switches(),
        "concurrency": {"deepseek": 1, "kimi": 5, "tongyi": 5, "baiduai": 5, "doubao": 1},
        "batch_count": 1,
        "idle_sleep": 30,
        "task_interval_min": 1,
        "task_interval_max": 3,
        "enable_task_commit": True,
        "location": "",
        "hot_reload": True,
        "once": False,
        "dry_run": False,
        "cdp": {"port": 9222},
        "browser": {"wenxin_share": False, "fallback_auth": False},
        "doubao": {
            "mode": "browser", "wait_timeout": 180, "share": True,
            "cookie_rotate_on_fail": True, "max_retry": 2, "use_cookie_pool": True,
            "chrome_path": r"D:\code\doubao-spider\ungoogled-chromium_142.0.7444.175\chrome.exe",
            "profile_dir": "doubao_profiles",
            "proxy": "",
            "captcha_server": "", "captcha_key": "", "captcha_timeout": 120,
        },
    }


def load_config(config_path: str | None = None) -> dict:
    path = config_path or os.path.join(BASE_DIR, DEFAULT_CONFIG_NAME)
    if not os.path.exists(path):
        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(default_config(), f, ensure_ascii=False, indent=2)
        print(f"[config] 已生成默认配置: {path}")
        print("[config] 请填写 spider_token 后重新运行")

    with open(path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    defs = default_config()
    for k, v in defs.items():
        if k not in cfg:
            cfg[k] = v
        elif k == "cookie_pool" and isinstance(v, dict):
            cfg["cookie_pool"] = {**cfg.get("cookie_pool", {}), **v}

    cfg["_config_path"] = os.path.abspath(path)
    cfg["batch_count"] = int(cfg.get("batch_count") or 1)
    cfg["idle_sleep"] = float(cfg.get("idle_sleep") or 30)
    cfg["task_interval_min"] = float(cfg.get("task_interval_min") or 1)
    cfg["task_interval_max"] = float(cfg.get("task_interval_max") or 3)
    cfg["spider_token"] = (cfg.get("spider_token") or "").strip()
    cfg["enable_task_commit"] = bool(cfg.get("enable_task_commit", True))
    cfg["once"] = bool(cfg.get("once", False))

    # 确保 deps 在 sys.path 中
    deps_dir = os.path.join(BASE_DIR, "deps")
    if deps_dir not in sys.path:
        sys.path.insert(0, deps_dir)
    if BASE_DIR not in sys.path:
        sys.path.insert(0, BASE_DIR)

    return cfg