diff --git a/src/bot/main.py b/src/bot/main.py index 3084834..183a531 100644 --- a/src/bot/main.py +++ b/src/bot/main.py @@ -7,8 +7,7 @@ from telegram.ext import ( ) 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.tools.orchestrator import Orchestrator from src.bot.states import chat_state, ChatMode logging.basicConfig( @@ -18,8 +17,7 @@ logging.basicConfig( logger = logging.getLogger(__name__) settings = get_settings() -tool_runner = ToolRunner() -memory = Memory() +orchestrator = Orchestrator() async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -31,6 +29,7 @@ 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) + current_tool = orchestrator.get_default_tool() help_text = ( f"Я {settings.bot_name}, ваш ИИ-ассистент.\n\n" "Доступные команды:\n" @@ -38,11 +37,13 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE): "/help - Показать эту справку\n" "/mode confirm - Режим с подтверждением\n" "/mode auto - Автономный режим\n" + "/use qwen|open - Выбрать инструмент\n" "/cancel - Отменить текущее действие\n" "/qwen <текст> - Задать вопрос qwen-code\n" "/open <текст> - Задать вопрос opencode\n" "/forget - Очистить историю чата\n\n" - f"Текущий режим: {'с подтверждением' if mode == ChatMode.CONFIRM else 'автономный'}" + f"Текущий режим: {'с подтверждением' if mode == ChatMode.CONFIRM else 'автономный'}\n" + f"Инструмент по умолчанию: {current_tool}" ) await update.message.reply_text(help_text) @@ -65,6 +66,21 @@ async def mode_command(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text("Использование: /mode confirm | auto") +async def use_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + if not context.args: + current = orchestrator.get_default_tool() + await update.message.reply_text(f"Текущий инструмент: {current}") + return + + tool = context.args[0].lower() + if tool in ["qwen", "open"]: + tool = "qwen" if tool == "qwen" else "opencode" + orchestrator.set_default_tool(tool) + await update.message.reply_text(f"Инструмент изменён на {tool}") + else: + await update.message.reply_text("Использование: /use qwen | open") + + async def cancel_command(update: Update, context: ContextTypes.DEFAULT_TYPE): chat_id = update.effective_chat.id @@ -94,29 +110,25 @@ async def confirm_callback(update: Update, context: ContextTypes.DEFAULT_TYPE): action_type = pending.get("type") if action_type == "tool": prompt = pending.get("prompt") - tool = pending.get("tool", "opencode") + tool = pending.get("tool") chat_state.set_waiting_confirmation(chat_id, False) - await execute_tool(query, context, tool, prompt) + await execute_tool_query(query, 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 +async def execute_tool_query(update, tool: str, prompt: str): + chat_id = update.message.chat.id if hasattr(update, 'message') else update.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) + result, success = await orchestrator.ask(prompt, chat_id, tool) text = result[:4096] if len(result) > 4096 else result - if hasattr(message, 'message'): - await message.message.reply_text(text) + if hasattr(update, 'message'): + await update.message.reply_text(text) else: - await message.reply_text(text) + await update.reply_text(text) async def qwen_command(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -149,7 +161,7 @@ async def qwen_command(update: Update, context: ContextTypes.DEFAULT_TYPE): ) else: await update.message.reply_text("Думаю...") - await execute_tool(update, context, "qwen", prompt) + await execute_tool_query(update, "qwen", prompt) async def open_command(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -182,17 +194,43 @@ async def open_command(update: Update, context: ContextTypes.DEFAULT_TYPE): ) else: await update.message.reply_text("Думаю...") - await execute_tool(update, context, "opencode", prompt) + await execute_tool_query(update, "opencode", prompt) async def forget_command(update: Update, context: ContextTypes.DEFAULT_TYPE): chat_id = update.effective_chat.id - memory.clear_chat(chat_id) + orchestrator.memory.clear_chat(chat_id) await update.message.reply_text("История чата очищена.") -async def echo(update: Update, context: ContextTypes.DEFAULT_TYPE): - await update.message.reply_text(update.message.text) +async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE): + prompt = update.message.text + chat_id = update.effective_chat.id + 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": orchestrator.get_default_tool(), + "prompt": prompt + }) + + await update.message.reply_text( + f"Выполнить запрос?\n\n{prompt[:200]}...", + reply_markup=reply_markup + ) + else: + await update.message.reply_text("Думаю...") + tool = orchestrator.get_default_tool() + await execute_tool_query(update, tool, prompt) def main(): @@ -208,12 +246,13 @@ 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("use", use_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)) + application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message)) logger.info("Бот запущен") application.run_polling(allowed_updates=Update.ALL_TYPES) diff --git a/src/tools/orchestrator.py b/src/tools/orchestrator.py new file mode 100644 index 0000000..fa709c0 --- /dev/null +++ b/src/tools/orchestrator.py @@ -0,0 +1,76 @@ +import logging +from typing import Tuple, Optional +from config.config import get_settings +from src.tools.tool_runner import ToolRunner +from src.memory.memory import Memory + +logger = logging.getLogger(__name__) +settings = get_settings() + +SYSTEM_PROMPT = """Ты Валера - дружелюбный, умный и полезный программист-ассистент. +Ты помогаешь пользователям с программированием, отвечаешь на вопросы, объясняешь код и помогаешь решать задачи. +Будь кратким, но информативным. Используй кодовые блоки для примеров.""" + + +class Orchestrator: + def __init__(self): + self.tool_runner = ToolRunner() + self.memory = Memory() + self.default_tool = settings.default_tool + self.tool_limits = { + "qwen": {"failed": 0, "max_failures": 3}, + "opencode": {"failed": 0, "max_failures": 3} + } + + def _check_rate_limit_error(self, result: str) -> bool: + rate_limit_keywords = [ + "rate limit", "quota", "превышен", "лимит", + "too many requests", "429", "daily limit" + ] + result_lower = result.lower() + return any(keyword in result_lower for keyword in rate_limit_keywords) + + def _build_prompt(self, user_prompt: str, chat_id: int) -> str: + context = self.memory.get_context_for_prompt(chat_id) + + full_prompt = f"{SYSTEM_PROMPT}\n\n" + if context: + full_prompt += f"История разговора:\n{context}\n\n" + full_prompt += f"Вопрос пользователя: {user_prompt}" + + return full_prompt + + async def ask(self, prompt: str, chat_id: int, tool: Optional[str] = None) -> Tuple[str, bool]: + selected_tool = tool or self.default_tool + + full_prompt = self._build_prompt(prompt, chat_id) + + result, success = await self.tool_runner.run_tool(selected_tool, full_prompt) + + if not success and self._check_rate_limit_error(result): + logger.warning(f"Лимит превышен для {selected_tool}, пробую другой инструмент") + + self.tool_limits[selected_tool]["failed"] += 1 + + if self.tool_limits[selected_tool]["failed"] >= self.tool_limits[selected_tool]["max_failures"]: + alt_tool = "opencode" if selected_tool == "qwen" else "qwen" + logger.info(f"Переключаюсь на {alt_tool}") + + result, success = await self.tool_runner.run_tool(alt_tool, full_prompt) + selected_tool = alt_tool + + if success: + self.tool_limits[selected_tool]["failed"] = 0 + + self.memory.add_message(chat_id, "user", prompt) + self.memory.add_message(chat_id, "assistant", result) + + return result, success + + def set_default_tool(self, tool: str): + if tool in ["qwen", "opencode"]: + self.default_tool = tool + logger.info(f"Инструмент по умолчанию изменён на {tool}") + + def get_default_tool(self) -> str: + return self.default_tool