157 lines
5.2 KiB
Python
157 lines
5.2 KiB
Python
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from .ai import AIClientError
|
||
from .speech import SpeechRecognitionError
|
||
|
||
|
||
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)
|
||
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 ""
|