268 lines
11 KiB
Python
268 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
AI Provider Manager - управление переключением между AI-провайдерами.
|
||
|
||
Поддерживаемые провайдеры:
|
||
- qwen: Qwen Code CLI (основной)
|
||
- gigachat: GigaChat API (Сбер)
|
||
|
||
Использует единый интерфейс BaseAIProvider для всех провайдеров.
|
||
"""
|
||
|
||
import logging
|
||
from typing import Optional, Dict, Any, Callable, List
|
||
from dataclasses import dataclass
|
||
from enum import Enum
|
||
|
||
from bot.base_ai_provider import BaseAIProvider, ProviderResponse
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AIProvider(Enum):
|
||
"""Доступные AI-провайдеры."""
|
||
QWEN = "qwen"
|
||
GIGACHAT = "gigachat"
|
||
|
||
|
||
@dataclass
|
||
class ProviderInfo:
|
||
"""Информация о провайдере."""
|
||
id: str
|
||
name: str
|
||
description: str
|
||
available: bool
|
||
is_active: bool
|
||
|
||
|
||
class AIProviderManager:
|
||
"""
|
||
Менеджер управления AI-провайдерами.
|
||
|
||
Позволяет переключаться между провайдерами и выполнять запросы
|
||
через активного провайдера с поддержкой инструментов.
|
||
"""
|
||
|
||
def __init__(self, qwen_manager=None):
|
||
self._qwen_manager = qwen_manager
|
||
self._provider_status: Dict[str, bool] = {}
|
||
self._providers: Dict[str, BaseAIProvider] = {}
|
||
self._tools_registry: Dict[str, Any] = {}
|
||
|
||
# Инициализируем провайдеров
|
||
self._init_providers()
|
||
|
||
# Проверяем доступность провайдеров при инициализации
|
||
self._check_provider_status()
|
||
|
||
def _init_providers(self):
|
||
"""Инициализировать AI-провайдеров."""
|
||
# Qwen Code Provider
|
||
if self._qwen_manager:
|
||
from bot.providers.qwen_provider import QwenCodeProvider
|
||
self._providers[AIProvider.QWEN.value] = QwenCodeProvider(self._qwen_manager)
|
||
logger.info("Qwen Code Provider инициализирован")
|
||
|
||
# GigaChat Provider - создаём новый экземпляр напрямую
|
||
from bot.providers.gigachat_provider import GigaChatProvider
|
||
self._providers[AIProvider.GIGACHAT.value] = GigaChatProvider()
|
||
logger.info("GigaChat Provider инициализирован")
|
||
|
||
def set_tools_registry(self, tools_registry: Dict[str, Any]):
|
||
"""Установить реестр инструментов для всех провайдеров."""
|
||
self._tools_registry = tools_registry
|
||
|
||
def get_provider(self, provider_id: str) -> Optional[BaseAIProvider]:
|
||
"""Получить экземпляр провайдера."""
|
||
return self._providers.get(provider_id)
|
||
|
||
def _check_provider_status(self):
|
||
"""Проверка доступности провайдеров."""
|
||
# Проверяем Qwen
|
||
self._provider_status[AIProvider.QWEN.value] = True # Qwen всегда доступен
|
||
|
||
# Проверяем GigaChat
|
||
gigachat_provider = self._providers.get(AIProvider.GIGACHAT.value)
|
||
if gigachat_provider:
|
||
self._provider_status[AIProvider.GIGACHAT.value] = gigachat_provider.is_available()
|
||
else:
|
||
self._provider_status[AIProvider.GIGACHAT.value] = False
|
||
|
||
def get_available_providers(self) -> List[str]:
|
||
"""Получить список доступных провайдеров."""
|
||
return [
|
||
provider_id
|
||
for provider_id, available in self._provider_status.items()
|
||
if available
|
||
]
|
||
|
||
def is_provider_available(self, provider_id: str) -> bool:
|
||
"""Проверить доступен ли провайдер."""
|
||
return self._provider_status.get(provider_id, False)
|
||
|
||
def get_provider_info(self, provider_id: str, is_active: bool = False) -> ProviderInfo:
|
||
"""Получить информацию о провайдере."""
|
||
providers = {
|
||
AIProvider.QWEN.value: ProviderInfo(
|
||
id=AIProvider.QWEN.value,
|
||
name="Qwen Code",
|
||
description="Alibaba Qwen Code CLI — мощный AI-ассистент с поддержкой инструментов",
|
||
available=self.is_provider_available(AIProvider.QWEN.value),
|
||
is_active=is_active
|
||
),
|
||
AIProvider.GIGACHAT.value: ProviderInfo(
|
||
id=AIProvider.GIGACHAT.value,
|
||
name="GigaChat",
|
||
description="Sber GigaChat API — российская AI-модель от Сбера",
|
||
available=self.is_provider_available(AIProvider.GIGACHAT.value),
|
||
is_active=is_active
|
||
)
|
||
}
|
||
return providers.get(provider_id)
|
||
|
||
def get_all_providers_info(self, active_provider_id: str) -> List[ProviderInfo]:
|
||
"""Получить информацию обо всех провайдерах."""
|
||
return [
|
||
self.get_provider_info(AIProvider.QWEN.value, AIProvider.QWEN.value == active_provider_id),
|
||
self.get_provider_info(AIProvider.GIGACHAT.value, AIProvider.GIGACHAT.value == active_provider_id)
|
||
]
|
||
|
||
def switch_provider(self, user_id: int, provider_id: str, state_manager) -> tuple[bool, str]:
|
||
"""
|
||
Переключить AI-провайдер для пользователя.
|
||
|
||
Args:
|
||
user_id: ID пользователя
|
||
provider_id: ID провайдера ("qwen" или "gigachat")
|
||
state_manager: Менеджер состояний для обновления состояния пользователя
|
||
|
||
Returns:
|
||
(success: bool, message: str)
|
||
"""
|
||
if not self.is_provider_available(provider_id):
|
||
return False, f"❌ Провайдер {provider_id} недоступен"
|
||
|
||
state = state_manager.get(user_id)
|
||
state.current_ai_provider = provider_id
|
||
|
||
provider_info = self.get_provider_info(provider_id)
|
||
|
||
logger.info(f"Пользователь {user_id} переключен на {provider_id}")
|
||
|
||
return True, f"✅ Переключен на {provider_info.name}"
|
||
|
||
def get_current_provider(self, state) -> str:
|
||
"""Получить текущего провайдера пользователя."""
|
||
return state.current_ai_provider
|
||
|
||
async def execute_request(
|
||
self,
|
||
provider_id: str,
|
||
user_id: int,
|
||
prompt: str,
|
||
system_prompt: Optional[str] = None,
|
||
on_output: Optional[Callable[[str], Any]] = None,
|
||
on_chunk: Optional[Callable[[str], Any]] = None,
|
||
on_event: Optional[Callable[[Any], Any]] = None,
|
||
context: Optional[List[Dict[str, str]]] = None,
|
||
use_tools: bool = True
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Выполнить запрос через указанного провайдера с поддержкой инструментов.
|
||
|
||
Args:
|
||
provider_id: ID провайдера
|
||
user_id: ID пользователя
|
||
prompt: Запрос
|
||
system_prompt: Системный промпт
|
||
on_output: Callback для вывода
|
||
on_chunk: Callback для потокового вывода
|
||
on_event: Callback для событий
|
||
context: История диалога
|
||
use_tools: Использовать ли инструменты
|
||
|
||
Returns:
|
||
Dict с результатом:
|
||
- success: bool
|
||
- content: str
|
||
- error: str (если ошибка)
|
||
- provider: str
|
||
- metadata: dict
|
||
"""
|
||
provider = self._providers.get(provider_id)
|
||
|
||
if not provider:
|
||
return {
|
||
"success": False,
|
||
"error": f"Провайдер {provider_id} не найден",
|
||
"provider": provider_id
|
||
}
|
||
|
||
try:
|
||
# Используем универсальный метод process_with_tools
|
||
response = await provider.process_with_tools(
|
||
prompt=prompt,
|
||
system_prompt=system_prompt,
|
||
context=context,
|
||
tools_registry=self._tools_registry if use_tools else None,
|
||
on_chunk=on_chunk,
|
||
user_id=user_id
|
||
)
|
||
|
||
if response.success:
|
||
# Получаем информацию о модели из metadata ответа
|
||
model_name = None
|
||
if response.message and response.message.metadata:
|
||
model_name = response.message.metadata.get("model")
|
||
|
||
return {
|
||
"success": True,
|
||
"content": response.message.content if response.message else "",
|
||
"provider": provider_id,
|
||
"metadata": {
|
||
"provider_name": response.provider_name,
|
||
"usage": response.usage,
|
||
"tool_calls": len(response.message.tool_calls) if response.message and response.message.tool_calls else 0,
|
||
"model": model_name # Добавляем модель
|
||
}
|
||
}
|
||
else:
|
||
return {
|
||
"success": False,
|
||
"error": response.error,
|
||
"provider": provider_id
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Ошибка выполнения запроса через {provider_id}: {e}")
|
||
return {
|
||
"success": False,
|
||
"error": str(e),
|
||
"provider": provider_id
|
||
}
|
||
|
||
|
||
# Глобальный менеджер (будет инициализирован в bot.py)
|
||
ai_provider_manager: Optional[AIProviderManager] = None
|
||
|
||
|
||
def init_ai_provider_manager(qwen_manager, tools_registry=None) -> AIProviderManager:
|
||
"""Инициализировать глобальный AIProviderManager."""
|
||
global ai_provider_manager
|
||
ai_provider_manager = AIProviderManager(qwen_manager)
|
||
|
||
# Устанавливаем реестр инструментов если предоставлен
|
||
if tools_registry:
|
||
ai_provider_manager.set_tools_registry(tools_registry)
|
||
|
||
logger.info(f"AIProviderManager инициализирован. Доступные провайдеры: {ai_provider_manager.get_available_providers()}")
|
||
return ai_provider_manager
|
||
|
||
|
||
def get_ai_provider_manager() -> AIProviderManager:
|
||
"""Получить глобальный AIProviderManager."""
|
||
global ai_provider_manager
|
||
if ai_provider_manager is None:
|
||
raise RuntimeError("AIProviderManager не инициализирован. Вызовите init_ai_provider_manager().")
|
||
return ai_provider_manager
|