src/cli.py - master - Anna_brv/ASR-Service

src/cli.py
262 строки7 KB
Anna
first_commit
30 апр 2026, 21:52
30 апр 2026, 21:52283885c
"""CLI для взаимодействия с ASR Service API."""

import argparse
import json
import sys
from pathlib import Path

from .client import ASRClient, transcribe_file


def main():
    """Точка входа CLI."""
    parser = argparse.ArgumentParser(
        description="ASR Service CLI — транскрибация аудио через API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Примеры использования:
  %(prog)s transcribe audio.mp3 --api-key your-key
  %(prog)s transcribe audio.mp3 --wait --diarize
  %(prog)s status <task-id> --api-key your-key
  %(prog)s models --api-key your-key
  %(prog)s health
        """,
    )

    parser.add_argument(
        "--base-url",
        default="http://localhost:8000",
        help="Базовый URL API (по умолчанию: http://localhost:8000)",
    )
    parser.add_argument(
        "--api-key",
        default=None,
        help="API-ключ для аутентификации",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=300.0,
        help="Таймаут запросов в секундах (по умолчанию: 300)",
    )

    subparsers = parser.add_subparsers(dest="command", help="Команда")

    # Команда transcribe
    transcribe_parser = subparsers.add_parser(
        "transcribe",
        help="Транскрибировать аудиофайл",
    )
    transcribe_parser.add_argument(
        "file",
        type=Path,
        help="Путь к аудиофайлу",
    )
    transcribe_parser.add_argument(
        "--model",
        default=None,
        help="Модель для транскрибации (по умолчанию: silero)",
    )
    transcribe_parser.add_argument(
        "--diarize",
        action="store_true",
        help="Включить диаризацию спикеров",
    )
    transcribe_parser.add_argument(
        "--language",
        default="ru",
        help="Язык аудио (по умолчанию: ru)",
    )
    transcribe_parser.add_argument(
        "--wait",
        action="store_true",
        help="Ждать завершения обработки",
    )
    transcribe_parser.add_argument(
        "--poll-interval",
        type=float,
        default=2.0,
        help="Интервал опроса статуса в секундах",
    )
    transcribe_parser.add_argument(
        "--output",
        "-o",
        type=Path,
        default=None,
        help="Сохранить результат в файл",
    )

    # Команда status
    status_parser = subparsers.add_parser(
        "status",
        help="Получить статус задачи",
    )
    status_parser.add_argument(
        "task_id",
        help="ID задачи",
    )

    # Команда cancel
    cancel_parser = subparsers.add_parser(
        "cancel",
        help="Отменить задачу",
    )
    cancel_parser.add_argument(
        "task_id",
        help="ID задачи",
    )

    # Команда restart
    restart_parser = subparsers.add_parser(
        "restart",
        help="Перезапустить задачу",
    )
    restart_parser.add_argument(
        "task_id",
        help="ID задачи",
    )

    # Команда list
    list_parser = subparsers.add_parser(
        "list",
        help="Список задач",
    )
    list_parser.add_argument(
        "--page",
        type=int,
        default=1,
        help="Номер страницы",
    )
    list_parser.add_argument(
        "--page-size",
        type=int,
        default=20,
        help="Размер страницы",
    )
    list_parser.add_argument(
        "--status",
        default=None,
        help="Фильтр по статусу",
    )

    # Команда models
    models_parser = subparsers.add_parser(
        "models",
        help="Список моделей",
    )

    # Команда load-model
    load_parser = subparsers.add_parser(
        "load-model",
        help="Загрузить модель",
    )
    load_parser.add_argument(
        "model_name",
        help="Название модели",
    )

    # Команда unload-model
    unload_parser = subparsers.add_parser(
        "unload-model",
        help="Выгрузить модель",
    )
    unload_parser.add_argument(
        "model_name",
        help="Название модели",
    )

    # Команда health
    subparsers.add_parser(
        "health",
        help="Проверка здоровья сервиса",
    )

    # Команда ping
    subparsers.add_parser(
        "ping",
        help="Проверка доступности API",
    )

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Создаём клиент
    client = ASRClient(
        base_url=args.base_url,
        api_key=args.api_key,
        timeout=args.timeout,
    )

    try:
        if args.command == "transcribe":
            result = client.transcribe(
                file_path=args.file,
                model_name=args.model,
                diarize=args.diarize,
                language=args.language,
                wait_for_result=args.wait,
                poll_interval=args.poll_interval,
            )

            if args.output:
                with open(args.output, "w", encoding="utf-8") as f:
                    json.dump(result, f, ensure_ascii=False, indent=2)
                print(f"Результат сохранён в {args.output}")
            else:
                print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "status":
            result = client.get_task_status(args.task_id)
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "cancel":
            result = client.cancel_task(args.task_id)
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "restart":
            result = client.restart_task(args.task_id)
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "list":
            result = client.list_tasks(
                page=args.page,
                page_size=args.page_size,
                status_filter=args.status,
            )
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "models":
            result = client.list_models()
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "load-model":
            result = client.load_model(args.model_name)
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "unload-model":
            result = client.unload_model(args.model_name)
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "health":
            result = client.health_check()
            print(json.dumps(result, ensure_ascii=False, indent=2))

        elif args.command == "ping":
            result = client.ping()
            print(json.dumps(result, ensure_ascii=False, indent=2))

    except FileNotFoundError as e:
        print(f"Ошибка: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Ошибка: {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        client.close()


if __name__ == "__main__":
    main()