diff --git a/src/bot/main.py b/src/bot/main.py index 18d6597..3084834 100644 --- a/src/bot/main.py +++ b/src/bot/main.py @@ -1,10 +1,15 @@ import asyncio import logging from telegram import Update -from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes +from telegram.ext import ( + Application, CommandHandler, MessageHandler, filters, + ContextTypes, CallbackQueryHandler +) +from telegram import InlineKeyboardButton, InlineKeyboardMarkup from config.config import get_settings from src.tools.tool_runner import ToolRunner from src.memory.memory import Memory +from src.bot.states import chat_state, ChatMode logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", @@ -25,18 +30,95 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + mode = chat_state.get_mode(update.effective_chat.id) help_text = ( f"Я {settings.bot_name}, ваш ИИ-ассистент.\n\n" "Доступные команды:\n" "/start - Начать работу\n" "/help - Показать эту справку\n" + "/mode confirm - Режим с подтверждением\n" + "/mode auto - Автономный режим\n" + "/cancel - Отменить текущее действие\n" "/qwen <текст> - Задать вопрос qwen-code\n" "/open <текст> - Задать вопрос opencode\n" - "/forget - Очистить историю чата\n" + "/forget - Очистить историю чата\n\n" + f"Текущий режим: {'с подтверждением' if mode == ChatMode.CONFIRM else 'автономный'}" ) await update.message.reply_text(help_text) +async def mode_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + if not context.args: + mode = chat_state.get_mode(update.effective_chat.id) + mode_name = "с подтверждением" if mode == ChatMode.CONFIRM else "автономный" + await update.message.reply_text(f"Текущий режим: {mode_name}") + return + + mode_arg = context.args[0].lower() + if mode_arg == "confirm": + chat_state.set_mode(update.effective_chat.id, ChatMode.CONFIRM) + await update.message.reply_text("Режим изменён: с подтверждением") + elif mode_arg == "auto": + chat_state.set_mode(update.effective_chat.id, ChatMode.AUTO) + await update.message.reply_text("Режим изменён: автономный") + else: + await update.message.reply_text("Использование: /mode confirm | auto") + + +async def cancel_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + chat_id = update.effective_chat.id + + if chat_state.is_waiting_confirmation(chat_id): + chat_state.set_waiting_confirmation(chat_id, False) + await update.message.reply_text("Ожидание подтверждения отменено.") + + task_id = chat_state.get_current_task(chat_id) + if task_id: + chat_state.set_current_task(chat_id, None) + await update.message.reply_text(f"Задача {task_id} отменена.") + else: + await update.message.reply_text("Нет активных задач для отмены.") + + +async def confirm_callback(update: Update, context: ContextTypes.DEFAULT_TYPE): + query = update.callback_query + await query.answer() + + chat_id = query.message.chat.id + data = query.data + + if data == "confirm_yes": + await query.edit_message_text("Подтверждено. Выполняю...") + pending = chat_state.get_pending_action(chat_id) + if pending: + action_type = pending.get("type") + if action_type == "tool": + prompt = pending.get("prompt") + tool = pending.get("tool", "opencode") + chat_state.set_waiting_confirmation(chat_id, False) + await execute_tool(query, context, tool, prompt) + elif data == "confirm_no": + chat_state.set_waiting_confirmation(chat_id, False) + await query.edit_message_text("Отменено.") + + +async def execute_tool(message, context, tool: str, prompt: str): + chat_id = message.message.chat.id if hasattr(message, 'message') else message.effective_chat.id + + memory.add_message(chat_id, "user", prompt) + + result, success = await tool_runner.run_tool(tool, prompt) + + memory.add_message(chat_id, "assistant", result) + + text = result[:4096] if len(result) > 4096 else result + + if hasattr(message, 'message'): + await message.message.reply_text(text) + else: + await message.reply_text(text) + + async def qwen_command(update: Update, context: ContextTypes.DEFAULT_TYPE): prompt = " ".join(context.args) if not prompt: @@ -44,13 +126,30 @@ async def qwen_command(update: Update, context: ContextTypes.DEFAULT_TYPE): return chat_id = update.effective_chat.id - memory.add_message(chat_id, "user", prompt) + mode = chat_state.get_mode(chat_id) - await update.message.reply_text("Думаю...") - result, success = await tool_runner.run_qwen(prompt) - - memory.add_message(chat_id, "assistant", result) - await update.message.reply_text(result[:4096] if len(result) > 4096 else result) + if mode == ChatMode.CONFIRM: + keyboard = [ + [ + InlineKeyboardButton("Да", callback_data="confirm_yes"), + InlineKeyboardButton("Нет", callback_data="confirm_no") + ] + ] + reply_markup = InlineKeyboardMarkup(keyboard) + + chat_state.set_waiting_confirmation(chat_id, True, { + "type": "tool", + "tool": "qwen", + "prompt": prompt + }) + + await update.message.reply_text( + f"Выполнить запрос к qwen-code?\n\n{prompt[:200]}...", + reply_markup=reply_markup + ) + else: + await update.message.reply_text("Думаю...") + await execute_tool(update, context, "qwen", prompt) async def open_command(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -60,13 +159,30 @@ async def open_command(update: Update, context: ContextTypes.DEFAULT_TYPE): return chat_id = update.effective_chat.id - memory.add_message(chat_id, "user", prompt) + mode = chat_state.get_mode(chat_id) - await update.message.reply_text("Думаю...") - result, success = await tool_runner.run_opencode(prompt) - - memory.add_message(chat_id, "assistant", result) - await update.message.reply_text(result[:4096] if len(result) > 4096 else result) + if mode == ChatMode.CONFIRM: + keyboard = [ + [ + InlineKeyboardButton("Да", callback_data="confirm_yes"), + InlineKeyboardButton("Нет", callback_data="confirm_no") + ] + ] + reply_markup = InlineKeyboardMarkup(keyboard) + + chat_state.set_waiting_confirmation(chat_id, True, { + "type": "tool", + "tool": "opencode", + "prompt": prompt + }) + + await update.message.reply_text( + f"Выполнить запрос к opencode?\n\n{prompt[:200]}...", + reply_markup=reply_markup + ) + else: + await update.message.reply_text("Думаю...") + await execute_tool(update, context, "opencode", prompt) async def forget_command(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -91,9 +207,12 @@ def main(): application.add_handler(CommandHandler("start", start)) application.add_handler(CommandHandler("help", help_command)) + application.add_handler(CommandHandler("mode", mode_command)) + application.add_handler(CommandHandler("cancel", cancel_command)) application.add_handler(CommandHandler("qwen", qwen_command)) application.add_handler(CommandHandler("open", open_command)) application.add_handler(CommandHandler("forget", forget_command)) + application.add_handler(CallbackQueryHandler(confirm_callback)) application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo)) logger.info("Бот запущен") diff --git a/src/bot/states.py b/src/bot/states.py new file mode 100644 index 0000000..88150ac --- /dev/null +++ b/src/bot/states.py @@ -0,0 +1,53 @@ +import logging +from typing import Dict, Optional +from enum import Enum + +logger = logging.getLogger(__name__) + + +class ChatMode(str, Enum): + CONFIRM = "confirm" + AUTO = "auto" + + +class ChatState: + def __init__(self): + self.states: Dict[int, dict] = {} + + def get_mode(self, chat_id: int) -> ChatMode: + return self.states.get(chat_id, {}).get("mode", ChatMode.CONFIRM) + + def set_mode(self, chat_id: int, mode: ChatMode): + if chat_id not in self.states: + self.states[chat_id] = {} + self.states[chat_id]["mode"] = mode + logger.info(f"Чат {chat_id} переключён в режим {mode}") + + def is_waiting_confirmation(self, chat_id: int) -> bool: + return self.states.get(chat_id, {}).get("waiting_confirmation", False) + + def set_waiting_confirmation(self, chat_id: int, waiting: bool, action_data: Optional[dict] = None): + if chat_id not in self.states: + self.states[chat_id] = {} + self.states[chat_id]["waiting_confirmation"] = waiting + if action_data: + self.states[chat_id]["pending_action"] = action_data + elif not waiting: + self.states[chat_id].pop("pending_action", None) + + def get_pending_action(self, chat_id: int) -> Optional[dict]: + return self.states.get(chat_id, {}).get("pending_action") + + def set_current_task(self, chat_id: int, task_id: Optional[str]): + if chat_id not in self.states: + self.states[chat_id] = {} + self.states[chat_id]["current_task"] = task_id + + def get_current_task(self, chat_id: int) -> Optional[str]: + return self.states.get(chat_id, {}).get("current_task") + + def clear(self, chat_id: int): + self.states.pop(chat_id, None) + + +chat_state = ChatState()