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 ""