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(), )