120 lines
4.7 KiB
Python
120 lines
4.7 KiB
Python
"""ИИ-агент для чата 'Планета Земля'."""
|
|
from typing import List, Optional
|
|
from uuid import UUID
|
|
|
|
from models.gigachat_types import GigaChatMessage
|
|
from prompts.persona import EARTH_PERSONA
|
|
|
|
from agents.gigachat_client import GigaChatClient
|
|
from services.cache_service import CacheService
|
|
|
|
|
|
class ChatAgent:
|
|
"""ИИ-агент для общения с детьми и родителями."""
|
|
|
|
def __init__(self, gigachat: GigaChatClient, cache: CacheService):
|
|
self.gigachat = gigachat
|
|
self.cache = cache
|
|
|
|
async def chat(
|
|
self,
|
|
user_id: UUID,
|
|
message: str,
|
|
conversation_id: Optional[str] = None,
|
|
model: str = "GigaChat-2",
|
|
) -> tuple[str, int]:
|
|
"""
|
|
Отправить сообщение и получить ответ.
|
|
|
|
Args:
|
|
user_id: ID пользователя
|
|
message: Текст сообщения
|
|
conversation_id: ID разговора (для контекста)
|
|
model: Модель GigaChat
|
|
|
|
Returns:
|
|
(ответ, количество использованных токенов)
|
|
"""
|
|
# Загружаем контекст из кэша
|
|
context_messages = []
|
|
if conversation_id:
|
|
cached_context = await self.cache.get_context(str(conversation_id))
|
|
# Фильтруем системные сообщения из кэша - они не должны там храниться
|
|
context_messages = [
|
|
GigaChatMessage(role=msg["role"], content=msg["content"])
|
|
for msg in cached_context
|
|
if msg["role"] != "system"
|
|
]
|
|
|
|
# Системное сообщение ВСЕГДА должно быть первым
|
|
system_message = GigaChatMessage(role="system", content=EARTH_PERSONA)
|
|
# Убеждаемся, что системное сообщение первое (удаляем все системные сообщения и добавляем одно в начало)
|
|
context_messages = [msg for msg in context_messages if msg.role != "system"]
|
|
context_messages.insert(0, system_message)
|
|
|
|
# Добавляем текущее сообщение пользователя
|
|
context_messages.append(GigaChatMessage(role="user", content=message))
|
|
|
|
# Отправляем запрос (не передаем message отдельно, т.к. оно уже в context_messages)
|
|
response = await self.gigachat.chat_with_response(
|
|
message="", # Пустое, т.к. сообщение уже добавлено в context_messages
|
|
context=context_messages,
|
|
model=model,
|
|
temperature=0.7,
|
|
max_tokens=1500,
|
|
)
|
|
|
|
assistant_message = response.choices[0].message.content
|
|
tokens_used = response.usage.total_tokens
|
|
|
|
# Сохраняем в контекст
|
|
if conversation_id:
|
|
await self.cache.add_message(str(conversation_id), "user", message)
|
|
await self.cache.add_message(str(conversation_id), "assistant", assistant_message)
|
|
|
|
return assistant_message, tokens_used
|
|
|
|
async def chat_with_context(
|
|
self,
|
|
user_id: UUID,
|
|
message: str,
|
|
context: Optional[List[dict]] = None,
|
|
model: str = "GigaChat-2",
|
|
) -> tuple[str, int]:
|
|
"""
|
|
Отправить сообщение с явным контекстом.
|
|
|
|
Args:
|
|
user_id: ID пользователя
|
|
message: Текст сообщения
|
|
context: Явный контекст разговора
|
|
model: Модель GigaChat
|
|
|
|
Returns:
|
|
(ответ, количество использованных токенов)
|
|
"""
|
|
context_messages = [GigaChatMessage(role="system", content=EARTH_PERSONA)]
|
|
|
|
if context:
|
|
for msg in context:
|
|
context_messages.append(
|
|
GigaChatMessage(role=msg["role"], content=msg["content"])
|
|
)
|
|
|
|
context_messages.append(GigaChatMessage(role="user", content=message))
|
|
|
|
# Отправляем запрос (не передаем message отдельно, т.к. оно уже в context_messages)
|
|
response = await self.gigachat.chat_with_response(
|
|
message="", # Пустое, т.к. сообщение уже добавлено в context_messages
|
|
context=context_messages,
|
|
model=model,
|
|
temperature=0.7,
|
|
max_tokens=1500,
|
|
)
|
|
|
|
assistant_message = response.choices[0].message.content
|
|
tokens_used = response.usage.total_tokens
|
|
|
|
return assistant_message, tokens_used
|
|
|