src/models_api.py - master - vladelo/startups

src/models_api.py
211 строк11 KB
vladelo
Make smaller logging
18 авг 2025, 17:21
18 авг 2025, 17:21cd68f3a
import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Dict, List
from openai import AsyncOpenAI

from src.config import Config
from src.excel_file_manager import ExcelFileManager

logger = logging.getLogger(__name__)
config = Config()


class ModelStrategy(ABC):
    """Абстрактный класс для стратегий взаимодействия с моделями."""

    @abstractmethod
    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Получает ответ от модели на основе списка сообщений."""
        ...


class ChatGPTStrategy(ModelStrategy):
    """Стратегия для взаимодействия с ChatGPT через OpenAI API."""

    def __init__(self) -> None:
        """Инициализирует клиент OpenAI с API-ключом из конфигурации."""
        self.client = AsyncOpenAI(api_key=config.OPENAI_API_KEY)
        # Используем более дешевые модели
        self.model = getattr(config, 'OPENAI_MODEL', 'gpt-4o-mini')  # или gpt-3.5-turbo
        logger.info(f'Инициализирована стратегия {self.__class__.__name__} с моделью {self.model}')

    def _format_messages_for_log(self, messages: list, max_len: int = 50) -> str:
        lines = [f'Отправка сообщений в модель (всего: {len(messages)}):']
        for idx, msg in enumerate(messages, 1):
            role = msg.get('role', 'unknown')
            content = msg.get('content', '')
            content_len = len(content)
            if content_len > max_len:
                display_content = content[:max_len].replace('\n', ' ').replace('\r', ' ') + '... (truncated)'
            else:
                display_content = content.replace('\n', ' ').replace('\r', ' ')
            lines.append(f'  {idx}. role={role}; len={content_len}; text="{display_content}"')
        return '\n'.join(lines)

    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Отправляет запрос к ChatGPT и возвращает ответ."""
        try:
            # logger.info(f'[{self.__class__.__name__}] Отправка запроса, модель: {self.model}')
            # logger.debug(f'[{self.__class__.__name__}] {self._format_messages_for_log(messages)}')

            # Убираем web_search_options - они поддерживаются только в gpt-4o
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                max_tokens=4000,
            )

            content = response.choices[0].message.content.strip()
            token_usage = response.usage.completion_tokens if hasattr(response, 'usage') else 'неизвестно'
            logger.info(
                f'[{self.__class__.__name__}] Получен ответ, длина: {len(content)} символов, использовано токенов: {token_usage}',
            )
            return content
        except Exception as e:
            logger.error(f'[{self.__class__.__name__}] Ошибка при запросе: {e}')
            raise ValueError(f'Не удалось получить ответ от ChatGPT: {e}')


class ChatGPTFileStrategy(ModelStrategy):
    """Стратегия для взаимодействия с ChatGPT через OpenAI API для обработки файлов."""

    def __init__(self) -> None:
        """Инициализирует клиент OpenAI с API-ключом из конфигурации."""
        self.client = AsyncOpenAI(api_key=config.OPENAI_API_KEY)
        # Используем более дешевую модель, которая поддерживает файлы
        self.model = getattr(config, 'OPENAI_FILE_MODEL', 'gpt-4o-mini')
        logger.info(f'Инициализирована стратегия {self.__class__.__name__} с моделью {self.model}')

    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Отправляет запрос к ChatGPT и возвращает ответ."""
        try:
            logger.info(
                f'[{self.__class__.__name__}] Отправка запроса к файловой модели: {self.model}',
            )
            logger.debug(f'[{self.__class__.__name__}] Сообщения: {messages}')

            response = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.5,
                max_tokens=4000,
            )

            content = response.choices[0].message.content.strip()
            token_usage = response.usage.completion_tokens if hasattr(response, 'usage') else 'неизвестно'
            logger.info(
                f'[{self.__class__.__name__}] Получен ответ от файловой модели, длина: {len(content)} символов, использовано токенов: {token_usage}',
            )
            return content
        except Exception as e:
            logger.error(f'[{self.__class__.__name__}] Ошибка при запросе к файловой модели: {e}')
            raise ValueError(f'Не удалось получить ответ от ChatGPT (файловая модель): {e}')


class ExcelSearchStrategy(ModelStrategy):
    """Стратегия для поиска в Excel файле с использованием OpenAI Assistants API."""

    def __init__(self) -> None:
        """Инициализирует клиент OpenAI с API-ключом из конфигурации."""
        self.client = AsyncOpenAI(api_key=config.OPENAI_API_KEY)
        self.model = getattr(config, 'OPENAI_FILE_MODEL', 'gpt-4o-mini')
        self._file_manager = ExcelFileManager()
        logger.info(f'Инициализирована стратегия {self.__class__.__name__} с моделью {self.model}')

    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Находит наиболее релевантную информацию для запроса пользователя."""
        try:
            user_query = messages[-1]['content']

            # Создаем assistant с file_search
            assistant = await self.client.beta.assistants.create(
                name="Excel Search Assistant",
                instructions=(
                    "Ты помощник для поиска информации в Excel файлах. "
                    "Найди в загруженном файле информацию, которая соответствует запросу пользователя. "
                    "Верни заголовок таблицы и минимум 3 строки, которые подходят под запрос. "
                    "Не добавляй пояснений, только данные из таблицы."
                ),
                model=self.model,
                tools=[{"type": "file_search"}],
                tool_resources={
                    "file_search": {
                        "vector_store_ids": [self._file_manager._vector_store_id]
                    }
                }
            )

            # Создаем thread
            thread = await self.client.beta.threads.create()

            # Добавляем сообщение
            await self.client.beta.threads.messages.create(
                thread_id=thread.id,
                role="user",
                content=f"Найди в Excel файле информацию по запросу: {user_query}"
            )

            # Запускаем assistant
            run = await self.client.beta.threads.runs.create(
                thread_id=thread.id,
                assistant_id=assistant.id
            )

            # Ждем завершения
            while run.status in ['queued', 'in_progress', 'cancelling']:
                await asyncio.sleep(1)
                run = await self.client.beta.threads.runs.retrieve(
                    thread_id=thread.id,
                    run_id=run.id
                )

            if run.status == 'completed':
                # Получаем ответ
                messages = await self.client.beta.threads.messages.list(
                    thread_id=thread.id
                )

                response_content = messages.data[0].content[0].text.value

                # Очищаем ресурсы
                await self.client.beta.threads.delete(thread.id)
                await self.client.beta.assistants.delete(assistant.id)

                logger.info(f'Получен ответ от Excel search, длина: {len(response_content)} символов')
                return response_content
            else:
                # Очищаем ресурсы в случае ошибки
                await self.client.beta.threads.delete(thread.id)
                await self.client.beta.assistants.delete(assistant.id)
                raise ValueError(f"Assistant run failed with status: {run.status}")

        except Exception as e:
            logger.error(f'Ошибка при поиске в Excel файле: {e}')
            raise ValueError(f'Не удалось найти релевантную информацию: {e}')


class ModelAPI:
    """Фасад для взаимодействия с различными моделями."""

    def __init__(self, strategy: ModelStrategy) -> None:
        """Инициализирует API с выбранной стратегией."""
        self._strategy = strategy
        logger.info(f'Инициализирован ModelAPI с стратегией {strategy.__class__.__name__}')

    @property
    def strategy(self) -> ModelStrategy:
        """Возвращает текущую стратегию."""
        return self._strategy

    @strategy.setter
    def strategy(self, strategy: ModelStrategy) -> None:
        """Устанавливает новую стратегию."""
        logger.info(f'Смена стратегии с {self._strategy.__class__.__name__} на {strategy.__class__.__name__}')
        self._strategy = strategy

    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Получает ответ от модели с использованием текущей стратегии."""
        # logger.info(f'Запрос ответа через стратегию {self._strategy.__class__.__name__}')
        return await self._strategy.get_response(messages)