118 lines
5.2 KiB
Python
118 lines
5.2 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Простой клиент для GigaChat API с поддержкой повторных попыток и circuit breaker.
|
||
Использует tenacity для retry и pybreaker для circuit breaker.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from typing import List, Dict, Optional
|
||
|
||
from gigachat import GigaChat
|
||
from gigachat.models import Chat, Messages, MessagesRole
|
||
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
|
||
import pybreaker
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Circuit breaker: при 3 ошибках подряд открывается на 60 секунд
|
||
breaker = pybreaker.CircuitBreaker(fail_max=3, reset_timeout=60)
|
||
|
||
|
||
class GigaClient:
|
||
def __init__(self, api_key: str, model: str, temperature: float, timeout: int, verify_ssl: bool = False):
|
||
self.api_key = api_key
|
||
self.model = model
|
||
self.temperature = temperature
|
||
self.timeout = timeout
|
||
self.verify_ssl = verify_ssl
|
||
self._client = None
|
||
|
||
def _get_client(self) -> GigaChat:
|
||
"""Ленивое создание клиента (создаётся при первом запросе)."""
|
||
if self._client is None:
|
||
self._client = GigaChat(
|
||
credentials=self.api_key,
|
||
model=self.model,
|
||
verify_ssl_certs=self.verify_ssl,
|
||
timeout=self.timeout
|
||
)
|
||
logger.debug("GigaChat client initialized")
|
||
return self._client
|
||
|
||
@retry(
|
||
stop=stop_after_attempt(3),
|
||
wait=wait_exponential(multiplier=1, min=1, max=8),
|
||
retry=retry_if_exception_type((TimeoutError, ConnectionError))
|
||
)
|
||
@breaker
|
||
async def chat(
|
||
self,
|
||
history: List[Dict[str, str]],
|
||
query: str,
|
||
system_prompt: Optional[str] = None,
|
||
file_id: Optional[str] = None,
|
||
temperature: Optional[float] = None,
|
||
) -> str:
|
||
"""
|
||
Отправляет запрос к GigaChat и возвращает ответ с повторными попытками.
|
||
|
||
Аргументы:
|
||
history: список предыдущих сообщений (каждое с ключами role и content).
|
||
query: текущий запрос пользователя.
|
||
system_prompt: системный промпт (если задан, добавляется первым сообщением).
|
||
file_id: идентификатор загруженного файла (для Vision).
|
||
|
||
Возвращает:
|
||
str: ответ модели или сообщение об ошибке.
|
||
"""
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append(Messages(role=MessagesRole.SYSTEM, content=system_prompt))
|
||
for msg in history[-10:]: # ограничим историю 10 сообщениями для экономии токенов
|
||
role = MessagesRole.USER if msg['role'] == 'user' else MessagesRole.ASSISTANT
|
||
messages.append(Messages(role=role, content=msg['content']))
|
||
messages.append(Messages(role=MessagesRole.USER, content=query))
|
||
|
||
attachments = [file_id] if file_id else None
|
||
actual_temp = temperature if temperature is not None else self.temperature
|
||
|
||
try:
|
||
# Запускаем синхронный вызов giga.chat в отдельном потоке
|
||
response = await asyncio.to_thread(
|
||
self._get_client().chat,
|
||
Chat(
|
||
messages=messages,
|
||
model=self.model,
|
||
temperature=actual_temp,
|
||
n=1,
|
||
attachments=attachments
|
||
)
|
||
)
|
||
return response.choices[0].message.content
|
||
except pybreaker.CircuitBreakerError:
|
||
logger.error("Circuit breaker открыт, GigaChat временно недоступен")
|
||
return "Извините, сервис ИИ временно недоступен. Попробуйте позже."
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при вызове GigaChat: {e}", exc_info=True)
|
||
raise
|
||
|
||
async def upload_file(self, file_path: str) -> Optional[str]:
|
||
"""
|
||
Загружает файл в GigaChat для Vision-анализа (изображения).
|
||
Возвращает file_id для использования в chat().
|
||
"""
|
||
try:
|
||
with open(file_path, "rb") as f:
|
||
uploaded = await asyncio.to_thread(self._get_client().upload_file, f, "general")
|
||
if hasattr(uploaded, "id_"):
|
||
return uploaded.id_
|
||
if hasattr(uploaded, "id"):
|
||
return uploaded.id
|
||
if hasattr(uploaded, "file_id"):
|
||
return uploaded.file_id
|
||
data = uploaded.dict() if hasattr(uploaded, "dict") else {}
|
||
return data.get("id_") or data.get("id") or data.get("file_id")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка загрузки файла в GigaChat: {e}", exc_info=True)
|
||
return None |