507 lines
19 KiB
Python
507 lines
19 KiB
Python
import aiogram.types as types
|
||
from aiogram import F, Router
|
||
from aiogram.filters import StateFilter
|
||
from aiogram.fsm.context import FSMContext
|
||
|
||
from create_bot import orm, tz
|
||
from keyboards.reply_keyboards import (
|
||
CATEGORY_BY_TEXT,
|
||
USER_TYPE_BY_TEXT,
|
||
USER_TYPE_LABEL_BY_CODE,
|
||
get_back_to_menu_kb,
|
||
get_categories_kb,
|
||
get_consultation_actions_ikb,
|
||
get_consultations_ikb,
|
||
get_delete_profile_kb,
|
||
get_main_menu_kb,
|
||
get_profile_kb,
|
||
get_user_types_kb,
|
||
)
|
||
from states.client_states import AskQuestionStates, MainStates, ProfileStates
|
||
from utils.rag_api import ask_rag_answer
|
||
from utils.text_tools import format_llm_answer_html, split_plain_text_chunks, to_html
|
||
|
||
from datetime import datetime, timezone
|
||
|
||
|
||
client_router = Router()
|
||
|
||
DAILY_CONSULTATION_LIMIT = 5
|
||
CONSULTATION_MESSAGE_LIMIT = 5
|
||
|
||
|
||
def get_today_start_utc() -> datetime:
|
||
now_local = datetime.now(tz)
|
||
start_local = now_local.replace(hour=0, minute=0, second=0, microsecond=0)
|
||
return start_local.astimezone(timezone.utc)
|
||
|
||
|
||
def build_start_text() -> str:
|
||
return (
|
||
"Здравствуйте. Я юридический ИИ-консультант по законам РФ.\n\n"
|
||
"Я могу помочь:\n"
|
||
"— разобраться в ситуации;\n"
|
||
"— найти применимые нормы закона;\n"
|
||
"— объяснить их простыми словами;\n"
|
||
"— дать базовый план действий.\n\n"
|
||
"Ответ носит информационный характер и не заменяет консультацию юриста."
|
||
)
|
||
|
||
|
||
def build_help_text() -> str:
|
||
return (
|
||
"ℹ️ Как пользоваться ботом\n\n"
|
||
"1. Нажмите «Задать вопрос».\n"
|
||
"2. Выберите категорию.\n"
|
||
"3. Опишите ситуацию.\n"
|
||
"4. Укажите регион РФ.\n"
|
||
"5. Получите ответ с найденными источниками.\n\n"
|
||
"Пример хорошего вопроса:\n"
|
||
"«Работодатель не выплатил зарплату за апрель. Работаю официально, Москва. Что делать?»\n\n"
|
||
"Бот не заменяет юриста и не гарантирует результат спора."
|
||
)
|
||
|
||
|
||
def build_profile_text(user) -> str:
|
||
return (
|
||
"👤 Профиль\n\n"
|
||
f"Страна: {to_html(user.country or 'Россия')}\n"
|
||
f"Регион: {to_html(user.region or 'не указан')}\n"
|
||
f"Тип пользователя: {to_html(USER_TYPE_LABEL_BY_CODE.get(user.user_type, 'Физлицо'))}"
|
||
)
|
||
|
||
|
||
def build_consultations_text(consultations: list) -> str:
|
||
lines = ["📚 Ваши консультации:\n"]
|
||
|
||
for index, consultation in enumerate(consultations, start=1):
|
||
title = consultation.title or consultation.category
|
||
lines.append(
|
||
f"{index}. {to_html(title)} — {consultation.updated_at.astimezone(tz).strftime('%d.%m.%Y')}"
|
||
)
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def build_consultation_text(consultation, messages: list) -> str:
|
||
lines = [
|
||
f"📚 Консультация #{consultation.id}",
|
||
"",
|
||
f"Категория: {to_html(consultation.category)}",
|
||
f"Регион: {to_html(consultation.region or 'не указан')}",
|
||
f"Статус: {to_html(consultation.status)}",
|
||
"",
|
||
]
|
||
|
||
if not messages:
|
||
lines.append("Сообщений пока нет.")
|
||
return "\n".join(lines)
|
||
|
||
for message in messages[-6:]:
|
||
if message.role == "user":
|
||
title = "👤 Вопрос"
|
||
elif message.role == "assistant":
|
||
title = "🤖 Ответ"
|
||
else:
|
||
title = "ℹ️ Сообщение"
|
||
|
||
lines.append(title)
|
||
if message.role == "assistant":
|
||
lines.append(format_llm_answer_html(message.content[:1800]))
|
||
else:
|
||
lines.append(to_html(message.content[:1800]))
|
||
lines.append("")
|
||
|
||
return "\n".join(lines).strip()
|
||
|
||
|
||
def build_fallback_error_text(error_text: str) -> str:
|
||
if "No reliable law chunks" in error_text:
|
||
return (
|
||
"Я не нашел в базе надежную норму по этому вопросу. "
|
||
"Попробуйте уточнить ситуацию, добавить детали или выбрать другую категорию."
|
||
)
|
||
|
||
return error_text
|
||
|
||
|
||
async def show_main_menu(message: types.Message, state: FSMContext) -> None:
|
||
await message.answer(build_start_text(), reply_markup=get_main_menu_kb())
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
async def process_question(
|
||
message: types.Message,
|
||
state: FSMContext,
|
||
*,
|
||
category: str,
|
||
question: str,
|
||
region: str,
|
||
consultation_id: int | None = None,
|
||
) -> None:
|
||
user = await orm.get_user(message.from_user.id)
|
||
if user is None:
|
||
await show_main_menu(message, state)
|
||
return
|
||
|
||
if consultation_id is None:
|
||
consultations_today = await orm.count_user_consultations_since(
|
||
user_id=message.from_user.id,
|
||
since=get_today_start_utc(),
|
||
)
|
||
if consultations_today >= DAILY_CONSULTATION_LIMIT:
|
||
await message.answer(
|
||
"На сегодня достигнут лимит: 5 новых консультаций. Попробуйте завтра.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
else:
|
||
message_count = await orm.count_user_messages_in_consultation(consultation_id)
|
||
if message_count >= CONSULTATION_MESSAGE_LIMIT:
|
||
await message.answer(
|
||
"В этой консультации уже достигнут лимит: 5 сообщений пользователя. "
|
||
"Создайте новую консультацию.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
|
||
history = []
|
||
if consultation_id is not None:
|
||
history = await orm.get_consultation_messages(consultation_id)
|
||
|
||
await state.set_state(AskQuestionStates.processing)
|
||
await message.answer("⏳ Ищу подходящие нормы закона и готовлю ответ...")
|
||
|
||
try:
|
||
response = await ask_rag_answer(
|
||
user_id=message.from_user.id,
|
||
consultation_id=consultation_id,
|
||
question=question,
|
||
category=category,
|
||
region=region,
|
||
user_type=user.user_type,
|
||
history=history,
|
||
)
|
||
except RuntimeError as exc:
|
||
await message.answer(
|
||
build_fallback_error_text(str(exc)),
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
|
||
if user.region != region:
|
||
await orm.set_user_region(message.from_user.id, region)
|
||
|
||
raw_answer_text = response.get("answer", "Не удалось сформировать ответ.")
|
||
answer_chunks = [
|
||
format_llm_answer_html(chunk)
|
||
for chunk in split_plain_text_chunks(raw_answer_text)
|
||
]
|
||
saved_consultation_id = response.get("consultation_id")
|
||
|
||
for index, answer_chunk in enumerate(answer_chunks):
|
||
reply_markup = None
|
||
if index == len(answer_chunks) - 1 and saved_consultation_id is not None:
|
||
reply_markup = get_consultation_actions_ikb(saved_consultation_id)
|
||
await message.answer(answer_chunk, reply_markup=reply_markup)
|
||
await message.answer(
|
||
"Можете задать новый вопрос или открыть историю консультаций.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(F.text == "↩️ В меню", StateFilter("*"))
|
||
async def return_to_main_menu(message: types.Message, state: FSMContext):
|
||
await show_main_menu(message, state)
|
||
|
||
|
||
@client_router.message(F.text == "ℹ️ Помощь", StateFilter("*"))
|
||
async def show_help(message: types.Message, state: FSMContext):
|
||
await message.answer(build_help_text(), reply_markup=get_main_menu_kb())
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(F.text == "⚖️ Задать вопрос", StateFilter("*"))
|
||
async def ask_question_entry(message: types.Message, state: FSMContext):
|
||
await message.answer(
|
||
"Выберите категорию вопроса:",
|
||
reply_markup=get_categories_kb(),
|
||
)
|
||
await state.set_state(AskQuestionStates.choosing_category)
|
||
|
||
|
||
@client_router.message(
|
||
F.text.in_(list(CATEGORY_BY_TEXT.keys())),
|
||
AskQuestionStates.choosing_category,
|
||
)
|
||
async def set_question_category(message: types.Message, state: FSMContext):
|
||
await state.update_data(
|
||
category_text=message.text,
|
||
)
|
||
await message.answer(
|
||
"Опишите ситуацию одним сообщением.\n\n"
|
||
"Постарайтесь указать:\n"
|
||
"— что произошло;\n"
|
||
"— когда произошло;\n"
|
||
"— с кем спор;\n"
|
||
"— чего вы хотите добиться.",
|
||
reply_markup=get_back_to_menu_kb(),
|
||
)
|
||
await state.set_state(AskQuestionStates.waiting_question)
|
||
|
||
|
||
@client_router.message(AskQuestionStates.waiting_question)
|
||
async def set_question_text(message: types.Message, state: FSMContext):
|
||
question = (message.text or "").strip()
|
||
if not question:
|
||
await message.answer("Опишите ситуацию текстом одним сообщением.")
|
||
return
|
||
|
||
user = await orm.get_user(message.from_user.id)
|
||
|
||
await state.update_data(question=question)
|
||
|
||
if user and user.region:
|
||
await message.answer(
|
||
"Укажите регион РФ, где произошла ситуация.\n\n"
|
||
f"Сейчас в профиле указан регион: {to_html(user.region)}\n"
|
||
"Если он подходит, отправьте символ: -",
|
||
reply_markup=get_back_to_menu_kb(),
|
||
)
|
||
else:
|
||
await message.answer(
|
||
"Укажите регион РФ, где произошла ситуация.\n\n"
|
||
"Например:\n"
|
||
"Москва\n"
|
||
"Санкт-Петербург\n"
|
||
"Краснодарский край\n"
|
||
"Республика Татарстан",
|
||
reply_markup=get_back_to_menu_kb(),
|
||
)
|
||
|
||
await state.set_state(AskQuestionStates.waiting_region)
|
||
|
||
|
||
@client_router.message(AskQuestionStates.waiting_region)
|
||
async def set_region_and_process(message: types.Message, state: FSMContext):
|
||
payload = await state.get_data()
|
||
user = await orm.get_user(message.from_user.id)
|
||
region = (message.text or "").strip()
|
||
|
||
if region == "-" and user and user.region:
|
||
region = user.region
|
||
|
||
if not region:
|
||
await message.answer("Укажите регион РФ текстом.")
|
||
return
|
||
|
||
await process_question(
|
||
message,
|
||
state,
|
||
category=payload["category_text"],
|
||
question=payload["question"],
|
||
region=region,
|
||
)
|
||
|
||
|
||
@client_router.message(F.text == "📚 Мои консультации", StateFilter("*"))
|
||
async def show_consultations(message: types.Message, state: FSMContext):
|
||
consultations = await orm.list_user_consultations(message.from_user.id)
|
||
|
||
if not consultations:
|
||
await message.answer(
|
||
"У вас пока нет консультаций.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
|
||
await message.answer(
|
||
build_consultations_text(consultations),
|
||
reply_markup=get_consultations_ikb(consultations),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.callback_query(F.data.startswith("consultation:open:"))
|
||
async def open_consultation(callback: types.CallbackQuery, state: FSMContext):
|
||
consultation_id = int(callback.data.rsplit(":", 1)[-1])
|
||
consultation = await orm.get_consultation(consultation_id, callback.from_user.id)
|
||
|
||
if consultation is None:
|
||
await callback.answer("Консультация не найдена.", show_alert=True)
|
||
return
|
||
|
||
messages = await orm.get_consultation_messages(consultation_id)
|
||
await callback.message.answer(
|
||
build_consultation_text(consultation, messages),
|
||
reply_markup=get_consultation_actions_ikb(consultation_id),
|
||
)
|
||
await callback.answer()
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.callback_query(F.data.startswith("consultation:continue:"))
|
||
async def continue_consultation(callback: types.CallbackQuery, state: FSMContext):
|
||
consultation_id = int(callback.data.rsplit(":", 1)[-1])
|
||
consultation = await orm.get_consultation(consultation_id, callback.from_user.id)
|
||
|
||
if consultation is None:
|
||
await callback.answer("Консультация не найдена.", show_alert=True)
|
||
return
|
||
|
||
await state.update_data(
|
||
consultation_id=consultation.id,
|
||
category_text=consultation.category,
|
||
region=consultation.region,
|
||
)
|
||
await callback.message.answer(
|
||
"Напишите следующий вопрос по этой консультации одним сообщением.",
|
||
reply_markup=get_back_to_menu_kb(),
|
||
)
|
||
await callback.answer()
|
||
await state.set_state(AskQuestionStates.waiting_continue_question)
|
||
|
||
|
||
@client_router.callback_query(F.data.startswith("consultation:delete:"))
|
||
async def delete_consultation(callback: types.CallbackQuery, state: FSMContext):
|
||
consultation_id = int(callback.data.rsplit(":", 1)[-1])
|
||
await orm.delete_consultation(consultation_id, callback.from_user.id)
|
||
consultations = await orm.list_user_consultations(callback.from_user.id)
|
||
|
||
await callback.answer("Консультация удалена.")
|
||
|
||
if not consultations:
|
||
await callback.message.answer(
|
||
"Консультаций больше не осталось.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
|
||
await callback.message.answer(
|
||
build_consultations_text(consultations),
|
||
reply_markup=get_consultations_ikb(consultations),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(AskQuestionStates.waiting_continue_question)
|
||
async def process_continue_question(message: types.Message, state: FSMContext):
|
||
question = (message.text or "").strip()
|
||
if not question:
|
||
await message.answer("Напишите вопрос текстом одним сообщением.")
|
||
return
|
||
|
||
payload = await state.get_data()
|
||
consultation_id = payload.get("consultation_id")
|
||
consultation = await orm.get_consultation(consultation_id, message.from_user.id)
|
||
|
||
if consultation is None:
|
||
await message.answer(
|
||
"Консультация не найдена. Начните новую.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
return
|
||
|
||
await process_question(
|
||
message,
|
||
state,
|
||
category=consultation.category,
|
||
question=question,
|
||
region=consultation.region or payload.get("region"),
|
||
consultation_id=consultation.id,
|
||
)
|
||
|
||
|
||
@client_router.message(F.text == "👤 Профиль", StateFilter("*"))
|
||
async def show_profile(message: types.Message, state: FSMContext):
|
||
user = await orm.get_user(message.from_user.id)
|
||
|
||
if user is None:
|
||
await show_main_menu(message, state)
|
||
return
|
||
|
||
await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(F.text == "🌍 Изменить регион", StateFilter("*"))
|
||
async def change_region_start(message: types.Message, state: FSMContext):
|
||
await message.answer(
|
||
"Напишите новый регион РФ одним сообщением.",
|
||
reply_markup=get_back_to_menu_kb(),
|
||
)
|
||
await state.set_state(ProfileStates.waiting_region)
|
||
|
||
|
||
@client_router.message(ProfileStates.waiting_region)
|
||
async def change_region_finish(message: types.Message, state: FSMContext):
|
||
region = (message.text or "").strip()
|
||
if not region:
|
||
await message.answer("Регион не должен быть пустым.")
|
||
return
|
||
|
||
await orm.set_user_region(message.from_user.id, region)
|
||
user = await orm.get_user(message.from_user.id)
|
||
await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(F.text == "👔 Изменить тип пользователя", StateFilter("*"))
|
||
async def change_user_type_start(message: types.Message, state: FSMContext):
|
||
await message.answer(
|
||
"Выберите тип пользователя:",
|
||
reply_markup=get_user_types_kb(),
|
||
)
|
||
await state.set_state(ProfileStates.waiting_user_type)
|
||
|
||
|
||
@client_router.message(
|
||
F.text.in_(list(USER_TYPE_BY_TEXT.keys())),
|
||
ProfileStates.waiting_user_type,
|
||
)
|
||
async def change_user_type_finish(message: types.Message, state: FSMContext):
|
||
await orm.set_user_type(message.from_user.id, USER_TYPE_BY_TEXT[message.text])
|
||
user = await orm.get_user(message.from_user.id)
|
||
await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(F.text == "🗑 Удалить мои данные", StateFilter("*"))
|
||
async def delete_profile_start(message: types.Message, state: FSMContext):
|
||
await message.answer(
|
||
"Удалить историю консультаций и сбросить профиль?",
|
||
reply_markup=get_delete_profile_kb(),
|
||
)
|
||
await state.set_state(ProfileStates.confirm_delete)
|
||
|
||
|
||
@client_router.message(F.text == "✅ Да, удалить", ProfileStates.confirm_delete)
|
||
async def delete_profile_finish(message: types.Message, state: FSMContext):
|
||
await orm.delete_user(message.from_user.id)
|
||
await orm.create_user(
|
||
user_id=message.from_user.id,
|
||
username="@" + message.from_user.username if message.from_user.username else None,
|
||
fullname=to_html(message.from_user.full_name),
|
||
register_date=datetime.now(timezone.utc),
|
||
)
|
||
await message.answer(
|
||
"Данные удалены. История очищена, профиль сброшен.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|
||
await state.set_state(MainStates.main)
|
||
|
||
|
||
@client_router.message(MainStates.main)
|
||
async def main_menu_fallback(message: types.Message):
|
||
await message.answer(
|
||
"Выберите действие в меню ниже.",
|
||
reply_markup=get_main_menu_kb(),
|
||
)
|