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

src/client.py
337 строк10 KB
Anna
first_commit
30 апр 2026, 21:52
30 апр 2026, 21:52283885c
"""Python-клиент для взаимодействия с ASR Service API."""

import time
from pathlib import Path
from typing import Optional

import httpx


class ASRClient:
    """Клиент для работы с ASR Service API."""

    def __init__(
        self,
        base_url: str = "http://localhost:8000",
        api_key: Optional[str] = None,
        timeout: float = 300.0,
    ):
        """
        Инициализация клиента.

        Args:
            base_url: Базовый URL API.
            api_key: API-ключ для аутентификации.
            timeout: Таймаут запросов в секундах.
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout

        self._client = httpx.Client(
            base_url=self.base_url,
            timeout=self.timeout,
            headers={"X-API-Key": api_key} if api_key else {},
        )

    def close(self):
        """Закрыть соединение."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def transcribe(
        self,
        file_path: str | Path,
        model_name: Optional[str] = None,
        diarize: bool = False,
        language: str = "ru",
        wait_for_result: bool = False,
        poll_interval: float = 2.0,
    ) -> dict:
        """
        Транскрибировать аудиофайл.

        Args:
            file_path: Путь к аудиофайлу.
            model_name: Модель для транскрибации (None = модель по умолчанию).
            diarize: Включить диаризацию спикеров.
            language: Язык аудио.
            wait_for_result: Ждать ли завершения обработки.
            poll_interval: Интервал опроса статуса (сек).

        Returns:
            Результат транскрибации.
        """
        file_path = Path(file_path)

        if not file_path.exists():
            raise FileNotFoundError(f"Файл не найден: {file_path}")

        # Отправляем файл на транскрибацию
        with open(file_path, "rb") as f:
            files = {"file": f}
            data = {
                "diarize": str(diarize).lower(),
                "language": language,
            }
            if model_name:
                data["model_name"] = model_name

            response = self._client.post(
                "/api/v1/transcribe/",
                files=files,
                data=data,
            )
            response.raise_for_status()

        task_info = response.json()
        task_id = task_info["task_id"]

        # Если нужно ждать результат
        if wait_for_result:
            return self.wait_for_task(task_id, poll_interval)

        return task_info

    def wait_for_task(self, task_id: str, poll_interval: float = 2.0) -> dict:
        """
        Дождаться завершения задачи.

        Args:
            task_id: ID задачи.
            poll_interval: Интервал опроса (сек).

        Returns:
            Результат задачи.

        Raises:
            TimeoutError: Если задача не завершена за время timeout.
        """
        start_time = time.time()

        while True:
            if time.time() - start_time > self.timeout:
                raise TimeoutError(f"Таймаут ожидания задачи {task_id}")

            task_info = self.get_task_status(task_id)
            status = task_info["status"]

            if status in ["completed", "failed", "cancelled"]:
                return task_info

            time.sleep(poll_interval)

    def get_task_status(self, task_id: str) -> dict:
        """
        Получить статус задачи.

        Args:
            task_id: ID задачи.

        Returns:
            Информация о задаче.
        """
        response = self._client.get(f"/api/v1/transcribe/{task_id}")
        response.raise_for_status()
        return response.json()

    def cancel_task(self, task_id: str) -> dict:
        """
        Отменить задачу.

        Args:
            task_id: ID задачи.

        Returns:
            Результат отмены.
        """
        response = self._client.post(f"/api/v1/tasks/{task_id}/cancel")
        response.raise_for_status()
        return response.json()

    def restart_task(self, task_id: str) -> dict:
        """
        Перезапустить задачу.

        Args:
            task_id: ID задачи.

        Returns:
            Информация о новой задаче.
        """
        response = self._client.post(f"/api/v1/tasks/{task_id}/restart")
        response.raise_for_status()
        return response.json()

    def list_tasks(
        self,
        page: int = 1,
        page_size: int = 20,
        status_filter: Optional[str] = None,
    ) -> dict:
        """
        Получить список задач.

        Args:
            page: Номер страницы.
            page_size: Размер страницы.
            status_filter: Фильтр по статусу.

        Returns:
            Список задач.
        """
        params = {"page": page, "page_size": page_size}
        if status_filter:
            params["status_filter"] = status_filter

        response = self._client.get("/api/v1/tasks", params=params)
        response.raise_for_status()
        return response.json()

    def clear_queue(self) -> dict:
        """
        Очистить очередь задач.

        Returns:
            Результат очистки.
        """
        response = self._client.delete("/api/v1/tasks/clear")
        response.raise_for_status()
        return response.json()

    def list_models(self) -> dict:
        """
        Получить список моделей.

        Returns:
            Список доступных моделей.
        """
        response = self._client.get("/api/v1/models")
        response.raise_for_status()
        return response.json()

    def load_model(self, model_name: str) -> dict:
        """
        Загрузить модель.

        Args:
            model_name: Название модели.

        Returns:
            Результат загрузки.
        """
        response = self._client.post(
            "/api/v1/models/load",
            json={"model_name": model_name},
        )
        response.raise_for_status()
        return response.json()

    def unload_model(self, model_name: str) -> dict:
        """
        Выгрузить модель.

        Args:
            model_name: Название модели.

        Returns:
            Результат выгрузки.
        """
        response = self._client.post(
            "/api/v1/models/unload",
            params={"model_name": model_name},
        )
        response.raise_for_status()
        return response.json()

    def get_model_status(self, model_name: str) -> dict:
        """
        Получить статус модели.

        Args:
            model_name: Название модели.

        Returns:
            Информация о модели.
        """
        response = self._client.get(f"/api/v1/models/{model_name}/status")
        response.raise_for_status()
        return response.json()

    def health_check(self) -> dict:
        """
        Проверить здоровье сервиса.

        Returns:
            Статус сервиса.
        """
        response = self._client.get("/api/v1/health")
        response.raise_for_status()
        return response.json()

    def ping(self) -> dict:
        """
        Проверить доступность API.

        Returns:
            Ответ сервера.
        """
        response = self._client.get("/api/v1/ping")
        response.raise_for_status()
        return response.json()


# Утилиты для удобного использования

def transcribe_file(
    file_path: str | Path,
    api_key: Optional[str] = None,
    base_url: str = "http://localhost:8000",
    model_name: Optional[str] = None,
    diarize: bool = False,
    wait: bool = True,
) -> dict:
    """
    Быстрая транскрибация файла.

    Args:
        file_path: Путь к файлу.
        api_key: API-ключ.
        base_url: URL API.
        model_name: Модель (None = по умолчанию).
        diarize: Диаризация спикеров.
        wait: Ждать ли результат.

    Returns:
        Результат транскрибации.
    """
    with ASRClient(base_url=base_url, api_key=api_key) as client:
        return client.transcribe(
            file_path=file_path,
            model_name=model_name,
            diarize=diarize,
            wait_for_result=wait,
        )


if __name__ == "__main__":
    # Пример использования
    import sys

    if len(sys.argv) < 2:
        print("Использование: python -m src.client <audio_file> [api_key]")
        sys.exit(1)

    audio_file = sys.argv[1]
    api_key = sys.argv[2] if len(sys.argv) > 2 else None

    result = transcribe_file(audio_file, api_key=api_key, wait=True)
    print(f"Статус: {result['status']}")

    if result["status"] == "completed" and "result" in result:
        print(f"Текст: {result['result']['text']}")