kg-agent/scripts/check_step0_connections.py

"""Run minimal DeepSeek and Redis connection checks for Step0."""

from __future__ import annotations

import argparse

from step0_pre_prepare.connection_checks import (
    ConnectionCheckError,
    Step0ConnectionSettings,
    check_deepseek_connection,
    check_redis_connection,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "service",
        choices=("deepseek", "redis", "all"),
        nargs="?",
        default="all",
        help="要测试的服务，默认测试全部",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    settings = Step0ConnectionSettings.from_env()
    checks = []
    if args.service in ("deepseek", "all"):
        checks.append(("DeepSeek", check_deepseek_connection))
    if args.service in ("redis", "all"):
        checks.append(("Redis", check_redis_connection))

    failed = False
    for label, check in checks:
        try:
            check(settings)
        except ConnectionCheckError as exc:
            failed = True
            print(f"[FAIL] {label}: {exc}")
        else:
            print(f"[OK] {label}: connection succeeded")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())