AstrBot/astrbot/__main__.py

import argparse
import asyncio
import mimetypes
import os
import sys
from pathlib import Path

import anyio

from astrbot.core import LogBroker, LogManager, db_helper, logger
from astrbot.core.config.default import VERSION
from astrbot.core.initial_loader import InitialLoader
from astrbot.core.utils.astrbot_path import (
    get_astrbot_config_path,
    get_astrbot_data_path,
    get_astrbot_knowledge_base_path,
    get_astrbot_plugin_path,
    get_astrbot_root,
    get_astrbot_site_packages_path,
    get_astrbot_skills_path,
    get_astrbot_temp_path,
)
from astrbot.core.utils.io import (
    download_dashboard,
    get_dashboard_version,
)


# 将父目录添加到 sys.path
sys.path.append(Path(__file__).parent.as_posix())

logo_tmpl = r"""
     ___           _______.___________..______      .______     ______   .___________.
    /   \         /       |           ||   _  \     |   _  \   /  __  \  |           |
   /  ^  \       |   (----`---|  |----`|  |_)  |    |  |_)  | |  |  |  | `---|  |----`
  /  /_\  \       \   \       |  |     |      /     |   _  <  |  |  |  |     |  |
 /  _____  \  .----)   |      |  |     |  |\  \----.|  |_)  | |  `--'  |     |  |
/__/     \__\ |_______/       |__|     | _| `._____||______/   \______/      |__|

"""


def check_env() -> None:
    # Python version check: require 3.12 or 3.13
    if not (sys.version_info.major == 3 and sys.version_info.minor in (12, 13)):
        sys.exit(1)

    astrbot_root = get_astrbot_root()
    if astrbot_root not in sys.path:
        sys.path.insert(0, astrbot_root)

    site_packages_path = get_astrbot_site_packages_path()
    if site_packages_path not in sys.path:
        sys.path.insert(0, site_packages_path)

    os.makedirs(get_astrbot_config_path(), exist_ok=True)
    os.makedirs(get_astrbot_plugin_path(), exist_ok=True)
    os.makedirs(get_astrbot_temp_path(), exist_ok=True)
    os.makedirs(get_astrbot_knowledge_base_path(), exist_ok=True)
    os.makedirs(get_astrbot_skills_path(), exist_ok=True)
    os.makedirs(site_packages_path, exist_ok=True)

    # 针对问题 #181 的临时解决方案
    mimetypes.add_type("text/javascript", ".js")
    mimetypes.add_type("text/javascript", ".mjs")
    mimetypes.add_type("application/json", ".json")


async def check_dashboard_files(webui_dir: str | None = None):
    """下载管理面板文件"""
    # 指定webui目录
    if webui_dir:
        if await anyio.Path(webui_dir).exists():
            logger.info(f"使用指定的 WebUI 目录: {webui_dir}")
            return webui_dir
        logger.warning(f"指定的 WebUI 目录 {webui_dir} 不存在,将使用默认逻辑｡")

    data_dist_path = os.path.join(get_astrbot_data_path(), "dist")
    if await anyio.Path(data_dist_path).exists():
        v = await get_dashboard_version()
        if v is not None:
            # 存在文件
            if v == f"v{VERSION}":
                logger.info("WebUI 版本已是最新｡")
            else:
                logger.warning(
                    f"检测到 WebUI 版本 ({v}) 与当前 AstrBot 版本 (v{VERSION}) 不符｡",
                )
        return data_dist_path

    logger.info(
        "开始下载管理面板文件...高峰期(晚上)可能导致较慢的速度｡如多次下载失败,请前往 https://github.com/AstrBotDevs/AstrBot/releases/latest 下载 dist.zip,并将其中的 dist 文件夹解压至 data 目录下｡",
    )

    try:
        await download_dashboard(version=f"v{VERSION}", latest=False)
    except Exception as e:
        logger.warning(
            f"下载指定版本(v{VERSION})的管理面板文件失败: {e},尝试下载最新版本｡",
        )
        try:
            await download_dashboard(latest=True)
        except Exception as e:
            logger.critical(f"下载管理面板文件失败: {e}｡")
            return None

    logger.info("管理面板下载完成｡")
    return data_dist_path


async def main_async(webui_dir_arg: str | None, log_broker: LogBroker) -> None:
    """主异步入口"""
    # 检查仪表板文件
    webui_dir = await check_dashboard_files(webui_dir_arg)
    if webui_dir is None:
        logger.warning(
            "管理面板文件检查失败,WebUI 功能将不可用｡"
            "请检查网络连接或手动指定 --webui-dir 参数｡",
        )

    db = db_helper

    # 打印 logo
    logger.info(logo_tmpl)

    core_lifecycle = InitialLoader(db, log_broker)
    core_lifecycle.webui_dir = webui_dir
    await core_lifecycle.start()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AstrBot")
    parser.add_argument(
        "--webui-dir",
        type=str,
        help="指定 WebUI 静态文件目录路径",
        default=None,
    )
    args = parser.parse_args()

    check_env()

    # 启动日志代理
    log_broker = LogBroker()
    LogManager.set_queue_handler(logger, log_broker)

    # 只使用一次 asyncio.run()
    asyncio.run(main_async(args.webui_dir, log_broker))