386 lines
14 KiB
Python
386 lines
14 KiB
Python
"""
|
||
GigaChat API Tool для Telegram CLI Bot
|
||
|
||
Инструмент для работы с GigaChat API (Сбер).
|
||
Поддерживает генерацию текста, чат-сессии и различные модели.
|
||
|
||
Документация: https://developers.sber.ru/docs/ru/gigachat
|
||
"""
|
||
|
||
import os
|
||
import base64
|
||
import httpx
|
||
import asyncio
|
||
from typing import Optional, List, Dict, Any
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime, timedelta
|
||
|
||
|
||
@dataclass
|
||
class GigaChatMessage:
|
||
"""Сообщение для чата с GigaChat"""
|
||
role: str # 'user', 'assistant', 'system'
|
||
content: str
|
||
|
||
|
||
@dataclass
|
||
class GigaChatConfig:
|
||
"""Конфигурация подключения к GigaChat API"""
|
||
client_id: str
|
||
client_secret: str
|
||
scope: str = "GIGACHAT_API_PERS"
|
||
auth_url: str = "https://ngw.devices.sberbank.ru:9443/api/v2/oauth"
|
||
model: str = "GigaChat-Pro"
|
||
api_url: str = "https://gigachat.devices.sberbank.ru/api/v2"
|
||
timeout: int = 60
|
||
|
||
|
||
class GigaChatTool:
|
||
"""
|
||
Инструмент для работы с GigaChat API
|
||
|
||
Пример использования:
|
||
config = GigaChatConfig(
|
||
client_id=os.getenv("GIGACHAT_CLIENT_ID"),
|
||
client_secret=os.getenv("GIGACHAT_CLIENT_SECRET"),
|
||
)
|
||
tool = GigaChatTool(config)
|
||
|
||
# Простой запрос
|
||
response = await tool.chat("Привет, как дела?")
|
||
|
||
# Чат с историей
|
||
messages = [
|
||
GigaChatMessage(role="system", content="Ты полезный ассистент."),
|
||
GigaChatMessage(role="user", content="Расскажи про Python"),
|
||
]
|
||
response = await tool.chat(messages=messages)
|
||
"""
|
||
|
||
def __init__(self, config: Optional[GigaChatConfig] = None):
|
||
self.config = config or self._load_config_from_env()
|
||
self._access_token: Optional[str] = None
|
||
self._token_expires: Optional[datetime] = None
|
||
self._chat_history: List[GigaChatMessage] = []
|
||
|
||
def _load_config_from_env(self) -> GigaChatConfig:
|
||
"""Загрузка конфигурации из переменных окружения"""
|
||
return GigaChatConfig(
|
||
client_id=os.getenv("GIGACHAT_CLIENT_ID", ""),
|
||
client_secret=os.getenv("GIGACHAT_CLIENT_SECRET", ""),
|
||
scope=os.getenv("GIGACHAT_SCOPE", "GIGACHAT_API_PERS"),
|
||
auth_url=os.getenv("GIGACHAT_AUTH_URL", "https://ngw.devices.sberbank.ru:9443/api/v2/oauth"),
|
||
model=os.getenv("GIGACHAT_MODEL", "GigaChat-Pro"),
|
||
)
|
||
|
||
def _get_auth_headers(self) -> Dict[str, str]:
|
||
"""Получение заголовков для авторизации"""
|
||
credentials = f"{self.config.client_id}:{self.config.client_secret}"
|
||
encoded_credentials = base64.b64encode(credentials.encode()).decode()
|
||
return {
|
||
"Authorization": f"Basic {encoded_credentials}",
|
||
"Content-Type": "application/x-www-form-urlencoded",
|
||
"RqUID": "00000000-0000-0000-0000-000000000000",
|
||
}
|
||
|
||
async def _get_access_token(self) -> str:
|
||
"""Получение access токена для API"""
|
||
# Проверяем кэш токена
|
||
if self._access_token and self._token_expires:
|
||
if datetime.now() < self._token_expires - timedelta(minutes=5):
|
||
return self._access_token
|
||
|
||
# Запрашиваем новый токен
|
||
async with httpx.AsyncClient(verify=False) as client:
|
||
response = await client.post(
|
||
self.config.auth_url,
|
||
headers=self._get_auth_headers(),
|
||
data={"scope": self.config.scope},
|
||
timeout=30,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
self._access_token = data["access_token"]
|
||
# Токен действителен 30 минут, кэшируем на 25 минут
|
||
self._token_expires = datetime.now() + timedelta(minutes=25)
|
||
|
||
return self._access_token
|
||
|
||
async def chat(
|
||
self,
|
||
messages: Optional[List[GigaChatMessage]] = None,
|
||
model: Optional[str] = None,
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 2000,
|
||
top_p: float = 0.1,
|
||
repetition_penalty: float = 1.0,
|
||
use_history: bool = True,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Отправка запроса к GigaChat API
|
||
|
||
Args:
|
||
messages: Список сообщений (если None, используется история чата)
|
||
model: Модель для генерации (если None, используется модель из конфига)
|
||
temperature: Температура генерации (0.0 - 2.0)
|
||
max_tokens: Максимальное количество токенов в ответе
|
||
top_p: Параметр top-p sampling
|
||
repetition_penalty: Штраф за повторения
|
||
use_history: Использовать ли историю чата
|
||
|
||
Returns:
|
||
Dict с ответом API:
|
||
- content: Текст ответа
|
||
- model: Использованная модель
|
||
- usage: Статистика использования токенов
|
||
- finish_reason: Причина завершения
|
||
"""
|
||
token = await self._get_access_token()
|
||
|
||
# Формируем сообщения
|
||
if messages is None:
|
||
if use_history:
|
||
messages = self._chat_history.copy()
|
||
else:
|
||
messages = []
|
||
elif use_history:
|
||
# Добавляем новые сообщения к истории
|
||
self._chat_history.extend(messages)
|
||
messages = self._chat_history.copy()
|
||
|
||
# Преобразуем сообщения в формат API
|
||
api_messages = [
|
||
{"role": msg.role, "content": msg.content}
|
||
for msg in messages
|
||
]
|
||
|
||
payload = {
|
||
"model": model or self.config.model,
|
||
"messages": api_messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"top_p": top_p,
|
||
"repetition_penalty": repetition_penalty,
|
||
}
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {token}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
async with httpx.AsyncClient(verify=False) as client:
|
||
response = await client.post(
|
||
f"{self.config.api_url}/chat/completions",
|
||
headers=headers,
|
||
json=payload,
|
||
timeout=self.config.timeout,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
# Добавляем ответ ассистента в историю
|
||
if use_history and data.get("choices"):
|
||
assistant_message = data["choices"][0]["message"]
|
||
self._chat_history.append(GigaChatMessage(
|
||
role=assistant_message["role"],
|
||
content=assistant_message["content"],
|
||
))
|
||
|
||
return {
|
||
"content": data["choices"][0]["message"]["content"] if data.get("choices") else "",
|
||
"model": data.get("model", self.config.model),
|
||
"usage": data.get("usage", {}),
|
||
"finish_reason": data["choices"][0]["finish_reason"] if data.get("choices") else "",
|
||
}
|
||
|
||
def clear_history(self):
|
||
"""Очистка истории чата"""
|
||
self._chat_history = []
|
||
|
||
def get_history(self) -> List[GigaChatMessage]:
|
||
"""Получение истории чата"""
|
||
return self._chat_history.copy()
|
||
|
||
def set_system_prompt(self, prompt: str):
|
||
"""Установка системного промпта (добавляется в начало истории)"""
|
||
# Удаляем старый системный промпт если есть
|
||
self._chat_history = [
|
||
msg for msg in self._chat_history if msg.role != "system"
|
||
]
|
||
# Добавляем новый в начало
|
||
self._chat_history.insert(0, GigaChatMessage(role="system", content=prompt))
|
||
|
||
async def generate_image(
|
||
self,
|
||
prompt: str,
|
||
model: str = " Kandinsky-2",
|
||
size: str = "1024x1024",
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Генерация изображений через GigaChat (Kandinsky)
|
||
|
||
Args:
|
||
prompt: Текстовое описание изображения
|
||
model: Модель для генерации
|
||
size: Размер изображения
|
||
|
||
Returns:
|
||
Dict с результатом генерации
|
||
"""
|
||
token = await self._get_access_token()
|
||
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"size": size,
|
||
}
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {token}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
async with httpx.AsyncClient(verify=False) as client:
|
||
# Запуск генерации
|
||
response = await client.post(
|
||
f"{self.config.api_url}/images/generations",
|
||
headers=headers,
|
||
json=payload,
|
||
timeout=self.config.timeout,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
return data
|
||
|
||
async def get_models(self) -> List[str]:
|
||
"""Получение списка доступных моделей"""
|
||
token = await self._get_access_token()
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {token}",
|
||
}
|
||
|
||
async with httpx.AsyncClient(verify=False) as client:
|
||
response = await client.get(
|
||
f"{self.config.api_url}/models",
|
||
headers=headers,
|
||
timeout=30,
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
return [model["id"] for model in data.get("data", [])]
|
||
|
||
|
||
# Утилита для создания инструмента в формате бота
|
||
def create_gigachat_tool():
|
||
"""
|
||
Создает экземпляр GigaChatTool с конфигурацией из окружения
|
||
|
||
Returns:
|
||
GigaChatTool или None если конфигурация не задана
|
||
"""
|
||
if not os.getenv("GIGACHAT_CLIENT_ID") or not os.getenv("GIGACHAT_CLIENT_SECRET"):
|
||
return None
|
||
|
||
return GigaChatTool()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# Пример использования
|
||
async def main():
|
||
tool = create_gigachat_tool()
|
||
if not tool:
|
||
print("GigaChat не настроен. Проверьте переменные окружения.")
|
||
return
|
||
|
||
# Простой запрос
|
||
response = await tool.chat("Привет! Расскажи кратко про себя.")
|
||
print(f"Ответ: {response['content']}")
|
||
print(f"Модель: {response['model']}")
|
||
print(f"Токены: {response['usage']}")
|
||
|
||
asyncio.run(main())
|
||
|
||
# ===========================================
|
||
# Интеграция с реестром инструментов бота
|
||
# ===========================================
|
||
|
||
from bot.tools import BaseTool, ToolResult, register_tool
|
||
|
||
|
||
@register_tool
|
||
class GigaChatCapability(BaseTool):
|
||
"""
|
||
Capability-обёртка для GigaChat API.
|
||
|
||
Позволяет использовать GigaChat через реестр инструментов бота.
|
||
"""
|
||
|
||
name = "gigachat"
|
||
description = "Генерация ответов AI через GigaChat API (Сбер). Альтернатива Qwen Code."
|
||
category = "ai"
|
||
|
||
def __init__(self):
|
||
self._provider = None
|
||
|
||
def _get_provider(self):
|
||
"""Ленивая инициализация провайдера"""
|
||
if self._provider is None:
|
||
from qwen_integration import GigaChatProvider
|
||
self._provider = GigaChatProvider()
|
||
return self._provider
|
||
|
||
async def execute(
|
||
self,
|
||
prompt: str,
|
||
system_prompt: Optional[str] = None,
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 2000,
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""
|
||
Выполнить запрос к GigaChat API.
|
||
|
||
Args:
|
||
prompt: Запрос пользователя
|
||
system_prompt: Системный промпт (роль ассистента)
|
||
temperature: Температура генерации (0.0-2.0)
|
||
max_tokens: Максимум токенов в ответе
|
||
"""
|
||
provider = self._get_provider()
|
||
|
||
if not provider.is_available():
|
||
return ToolResult(
|
||
success=False,
|
||
error=provider.get_error() or "GigaChat не доступен",
|
||
)
|
||
|
||
result = await provider.chat(
|
||
prompt=prompt,
|
||
system_prompt=system_prompt,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
)
|
||
|
||
if result.get("success"):
|
||
return ToolResult(
|
||
success=True,
|
||
data={
|
||
"content": result.get("content", ""),
|
||
"model": result.get("model", "GigaChat-Pro"),
|
||
"usage": result.get("usage", {}),
|
||
},
|
||
metadata={
|
||
"model": result.get("model"),
|
||
"tokens": result.get("usage"),
|
||
}
|
||
)
|
||
else:
|
||
return ToolResult(
|
||
success=False,
|
||
error=result.get("error", "Неизвестная ошибка GigaChat"),
|
||
)
|