fckbot/rag/services/giga_client.py

# -*- coding: utf-8 -*-
"""
Простой клиент для GigaChat API.
Только вызовы chat и upload_file, без бизнес-логики.
Используется AIService и другими модулями.
"""

import asyncio
import logging
from typing import List, Dict, Optional

from gigachat import GigaChat
from gigachat.models import Chat, Messages, MessagesRole

logger = logging.getLogger(__name__)


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

    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 Exception as e:
            logger.error(f"Ошибка при вызове GigaChat: {e}", exc_info=True)
            return "Извините, произошла ошибка при обращении к сервису ИИ."

    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")
                # В разных версиях SDK поле может называться id_, id или file_id
                if hasattr(uploaded, "id_"):
                    return uploaded.id_
                if hasattr(uploaded, "id"):
                    return uploaded.id
                if hasattr(uploaded, "file_id"):
                    return uploaded.file_id
                # fallback: пробуем преобразовать в словарь
                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