kandrusyak_bot/assistant_bot/yandex_ai.py

import logging
from pathlib import Path
from typing import Any

from .ai import AIClientError
from .speech import SpeechRecognitionError


logger = logging.getLogger(__name__)


class YandexAIError(AIClientError):
    pass


def describe_yandex_error(exc: Exception) -> str:
    response = getattr(exc, "response", None)
    status_code = getattr(response, "status_code", None)
    response_text = getattr(response, "text", None)
    if status_code and isinstance(response_text, str) and response_text.strip():
        return f"HTTP {status_code}: {response_text.strip()[:800]}"
    return str(exc)


def prepare_chat_messages(
    messages: list[dict[str, str]],
) -> list[dict[str, str]]:
    """Convert messages to the format accepted by Yandex prompt templates."""
    system_parts: list[str] = []
    chat_messages: list[dict[str, str]] = []

    for message in messages:
        role = message.get("role", "user")
        content = message.get("content", "")
        if role == "system":
            system_parts.append(content)
        else:
            chat_messages.append({"role": role, "content": content})

    if system_parts:
        chat_messages.insert(
            0,
            {"role": "system", "content": "\n\n".join(system_parts)},
        )

    return chat_messages


def create_async_sdk(folder_id: str) -> Any:
    try:
        from yandex_ai_studio_sdk import AsyncAIStudio
    except ImportError as exc:
        raise RuntimeError(
            "Yandex AI Studio SDK is not installed. "
            "Run: python -m pip install -r requirements.txt"
        ) from exc
    return AsyncAIStudio(folder_id=folder_id)


def create_sync_sdk(folder_id: str) -> Any:
    try:
        from yandex_ai_studio_sdk import AIStudio
    except ImportError as exc:
        raise RuntimeError(
            "Yandex AI Studio SDK is not installed. "
            "Run: python -m pip install -r requirements.txt"
        ) from exc
    return AIStudio(folder_id=folder_id)


class YandexAIClient:
    provider = "yandex"
    display_name = "Yandex AI Studio"

    def __init__(self, folder_id: str, sdk: Any | None = None) -> None:
        self.folder_id = folder_id.strip("/")
        self._sdk = sdk if sdk is not None else create_async_sdk(self.folder_id)

    def normalize_model(self, model: str) -> str:
        normalized = model.strip()
        if normalized.startswith("gpt://"):
            model_path = normalized.removeprefix("gpt://").strip("/")
            parts = model_path.split("/")
            if (
                not normalized.startswith(f"gpt://{self.folder_id}/")
                and len(parts) == 2
                and parts[-1] in {"latest", "rc"}
            ):
                return f"gpt://{self.folder_id}/{model_path}"
            return normalized
        return f"gpt://{self.folder_id}/{normalized.lstrip('/')}"

    async def list_models(self) -> list[str]:
        try:
            models = await self._sdk.chat.completions.list()
        except Exception as exc:
            raise YandexAIError(
                f"Не удалось получить список моделей: {describe_yandex_error(exc)}"
            ) from exc

        model_names = {
            str(uri)
            for model in models
            if (uri := getattr(model, "uri", None))
        }
        return sorted(model_names)

    async def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        json_mode: bool = False,
    ) -> str:
        sdk_messages = prepare_chat_messages(messages)

        try:
            completion = self._sdk.chat.completions(self.normalize_model(model))
            if json_mode:
                completion = completion.configure(response_format="json")
            result = await completion.run(sdk_messages, timeout=180)
            usage = getattr(result, "usage", None)
            if usage is not None:
                logger.info(
                    "AI usage provider=yandex model=%s input_tokens=%s "
                    "output_tokens=%s total_tokens=%s",
                    self.normalize_model(model),
                    getattr(
                        usage,
                        "prompt_tokens",
                        getattr(usage, "input_text_tokens", None),
                    ),
                    getattr(usage, "completion_tokens", None),
                    getattr(usage, "total_tokens", None),
                )
            content = getattr(result, "text", None)
            if content is None:
                content = getattr(result[0], "text", None)
        except Exception as exc:
            raise YandexAIError(
                f"Запрос к модели завершился ошибкой: {describe_yandex_error(exc)}"
            ) from exc

        if isinstance(content, str) and content.strip():
            return content.strip()
        raise YandexAIError("Yandex AI Studio вернула пустой ответ.")


class YandexSpeechRecognizer:
    def __init__(
        self,
        folder_id: str,
        language: str,
        model: str,
        sdk: Any | None = None,
    ) -> None:
        self.folder_id = folder_id.strip("/")
        self.language = language
        self.model = model
        self._sdk = sdk if sdk is not None else create_sync_sdk(self.folder_id)

    def transcribe(self, audio_path: str | Path) -> str:
        try:
            audio = Path(audio_path).read_bytes()
            recognizer = self._sdk.speechkit.speech_to_text(
                audio_format=self._sdk.speechkit.AudioFormat.OGG_OPUS,
                language_codes=self.language,
                model=self.model,
            )
            result = recognizer.run(audio, timeout=180)
            text = getattr(result, "text", None)
        except Exception as exc:
            raise SpeechRecognitionError("Yandex SpeechKit transcription failed") from exc

        return text.strip() if isinstance(text, str) else ""