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
|
||
|