Этап 4: Режимы подтверждения и прерывание
This commit is contained in:
mirivlad
parent
6c2f17e37a
commit
4f96c75b6d
143
src/bot/main.py
143
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)
|
||||
|
||||
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("Думаю...")
|
||||
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)
|
||||
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)
|
||||
|
||||
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("Думаю...")
|
||||
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)
|
||||
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("Бот запущен")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue