tg_bot_language/services/ai_model_service.py

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession
from database.models import AIModel, AIProvider, User
from typing import Optional, List, Tuple


# Дефолтная модель если в БД ничего нет
DEFAULT_MODEL = "gpt-4o-mini"
DEFAULT_PROVIDER = AIProvider.openai


class AIModelService:
    """Сервис для работы с AI моделями"""

    @staticmethod
    async def get_active_model(session: AsyncSession) -> Optional[AIModel]:
        """
        Получить активную AI модель

        Returns:
            AIModel или None если нет активной модели
        """
        result = await session.execute(
            select(AIModel).where(AIModel.is_active == True)
        )
        return result.scalar_one_or_none()

    @staticmethod
    async def get_active_model_name(session: AsyncSession) -> str:
        """
        Получить название активной модели

        Returns:
            Название модели (например "gpt-4o-mini") или дефолтное
        """
        model = await AIModelService.get_active_model(session)
        if model:
            return model.model_name
        return DEFAULT_MODEL

    @staticmethod
    async def get_active_provider(session: AsyncSession) -> AIProvider:
        """
        Получить провайдера активной модели

        Returns:
            AIProvider (OPENAI или GOOGLE)
        """
        model = await AIModelService.get_active_model(session)
        if model:
            return model.provider
        return DEFAULT_PROVIDER

    @staticmethod
    async def get_all_models(session: AsyncSession) -> List[AIModel]:
        """
        Получить все доступные модели

        Returns:
            Список всех моделей
        """
        result = await session.execute(
            select(AIModel).order_by(AIModel.provider, AIModel.model_name)
        )
        return list(result.scalars().all())

    @staticmethod
    async def set_active_model(session: AsyncSession, model_id: int) -> bool:
        """
        Установить активную модель по ID

        Args:
            model_id: ID модели для активации

        Returns:
            True если успешно, False если модель не найдена
        """
        # Проверяем существование модели
        result = await session.execute(
            select(AIModel).where(AIModel.id == model_id)
        )
        model = result.scalar_one_or_none()

        if not model:
            return False

        # Деактивируем все модели
        await session.execute(
            update(AIModel).values(is_active=False)
        )

        # Активируем выбранную
        model.is_active = True
        await session.commit()
        return True

    @staticmethod
    async def set_active_model_by_name(session: AsyncSession, model_name: str) -> bool:
        """
        Установить активную модель по названию

        Args:
            model_name: Название модели (например "gpt-4o-mini")

        Returns:
            True если успешно, False если модель не найдена
        """
        result = await session.execute(
            select(AIModel).where(AIModel.model_name == model_name)
        )
        model = result.scalar_one_or_none()

        if not model:
            return False

        # Деактивируем все модели
        await session.execute(
            update(AIModel).values(is_active=False)
        )

        # Активируем выбранную
        model.is_active = True
        await session.commit()
        return True

    @staticmethod
    async def create_model(
        session: AsyncSession,
        provider: AIProvider,
        model_name: str,
        display_name: str,
        is_active: bool = False
    ) -> AIModel:
        """
        Создать новую модель

        Args:
            provider: Провайдер (OPENAI, GOOGLE)
            model_name: Техническое название модели
            display_name: Отображаемое название
            is_active: Активна ли модель

        Returns:
            Созданная модель
        """
        # Если активируем новую модель, деактивируем остальные
        if is_active:
            await session.execute(
                update(AIModel).values(is_active=False)
            )

        model = AIModel(
            provider=provider,
            model_name=model_name,
            display_name=display_name,
            is_active=is_active
        )
        session.add(model)
        await session.commit()
        await session.refresh(model)
        return model

    @staticmethod
    async def ensure_default_models(session: AsyncSession):
        """
        Создать дефолтные модели если их нет в БД
        """
        result = await session.execute(select(AIModel))
        existing = list(result.scalars().all())

        if existing:
            return  # Модели уже есть

        # Создаём дефолтные модели
        default_models = [
            (AIProvider.openai, "gpt-4o-mini", "GPT-4o Mini", True),
            (AIProvider.openai, "gpt-5-nano", "GPT-5 Nano", False),
            (AIProvider.google, "gemini-2.5-flash-lite", "Gemini 2.5 Flash Lite", False),
        ]

        for provider, name, display, active in default_models:
            model = AIModel(
                provider=provider,
                model_name=name,
                display_name=display,
                is_active=active
            )
            session.add(model)

        await session.commit()

    @staticmethod
    async def get_user_model(session: AsyncSession, user_id: int) -> Optional[AIModel]:
        """
        Получить AI модель пользователя.
        Если у пользователя не выбрана модель, возвращает глобальную активную.

        Args:
            user_id: ID пользователя в БД

        Returns:
            AIModel или None
        """
        # Получаем пользователя
        result = await session.execute(
            select(User).where(User.id == user_id)
        )
        user = result.scalar_one_or_none()

        if user and user.ai_model_id:
            # У пользователя выбрана своя модель
            model_result = await session.execute(
                select(AIModel).where(AIModel.id == user.ai_model_id)
            )
            model = model_result.scalar_one_or_none()
            if model:
                return model

        # Fallback на глобальную активную модель
        return await AIModelService.get_active_model(session)

    @staticmethod
    async def get_user_model_info(session: AsyncSession, user_id: int) -> Tuple[str, AIProvider]:
        """
        Получить название модели и провайдера для пользователя.

        Args:
            user_id: ID пользователя в БД

        Returns:
            Tuple[model_name, provider]
        """
        model = await AIModelService.get_user_model(session, user_id)
        if model:
            return model.model_name, model.provider
        return DEFAULT_MODEL, DEFAULT_PROVIDER

    @staticmethod
    async def set_user_model(session: AsyncSession, user_id: int, model_id: Optional[int]) -> bool:
        """
        Установить AI модель для пользователя.

        Args:
            user_id: ID пользователя в БД
            model_id: ID модели или None для сброса на глобальную

        Returns:
            True если успешно
        """
        result = await session.execute(
            select(User).where(User.id == user_id)
        )
        user = result.scalar_one_or_none()

        if not user:
            return False

        # Проверяем существование модели если указан ID
        if model_id is not None:
            model_result = await session.execute(
                select(AIModel).where(AIModel.id == model_id)
            )
            if not model_result.scalar_one_or_none():
                return False

        user.ai_model_id = model_id
        await session.commit()
        return True