From 5779dd7b14d7890b6719d803da6bb18f665cb799 Mon Sep 17 00:00:00 2001 From: mirivlad Date: Wed, 18 Mar 2026 19:34:58 +0800 Subject: [PATCH] =?UTF-8?q?Add:=20=D0=BA=D0=BE=D0=BC=D0=B0=D0=BD=D0=B4?= =?UTF-8?q?=D0=B0=20/xray=20=D0=B4=D0=BB=D1=8F=20=D0=B4=D0=BE=D0=B1=D0=B0?= =?UTF-8?q?=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D1=8F=20=D0=BF=D0=BE=D0=BB=D1=8C?= =?UTF-8?q?=D0=B7=D0=BE=D0=B2=D0=B0=D1=82=D0=B5=D0=BB=D0=B5=D0=B9=20XRay,?= =?UTF-8?q?=20=D1=83=D0=BC=D0=B5=D0=BD=D1=8C=D1=88=D0=B5=D0=BD=D0=BE=20?= =?UTF-8?q?=D0=BB=D0=BE=D0=B3=D0=B8=D1=80=D0=BE=D0=B2=D0=B0=D0=BD=D0=B8?= =?UTF-8?q?=D0=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Добавлена команда /xray для управления пользователями XRay через SSH - SSH подключение к серверу mt.mirv.top с выполнением скрипта add_xray_user.sh - Генерация и отправка QR-кода для подключения к VPN - Интеграция с приложением Hiddify - Добавлены зависимости asyncssh и qrcode[pil] - Уменьшен уровень логирования (только WARNING и ERROR) Co-authored-by: Qwen-Coder --- .env.example | 6 + config/config.py | 10 +- requirements.txt | 2 + src/bot/main.py | 300 ++++++++++++++++++++++++++++++------- src/bot/states.py | 34 +++-- src/memory/memory.py | 9 +- src/scheduler/scheduler.py | 12 +- src/speech/speech.py | 18 +-- src/tools/orchestrator.py | 68 ++++++--- src/tools/tool_runner.py | 36 +++-- src/tools/xray.py | 172 +++++++++++++++++++++ 11 files changed, 539 insertions(+), 128 deletions(-) create mode 100644 src/tools/xray.py diff --git a/.env.example b/.env.example index 870e55a..b1ad266 100644 --- a/.env.example +++ b/.env.example @@ -32,3 +32,9 @@ YANDEX_FOLDER_ID= # Database DATABASE_URL=sqlite+aiosqlite:///./valera.db + +# XRay (mt.mirv.top) +XRAY_SSH_HOST=mt.mirv.top +XRAY_SSH_USER=root +XRAY_SSH_PASSWORD= +XRAY_ADD_USER_SCRIPT=/root/bin/add_xray_user.sh diff --git a/config/config.py b/config/config.py index 036e040..57d720d 100644 --- a/config/config.py +++ b/config/config.py @@ -28,9 +28,15 @@ class Settings(BaseSettings): gigachat_credentials: Optional[str] = None yandex_api_key: Optional[str] = None yandex_folder_id: Optional[str] = None - + database_url: str = "sqlite+aiosqlite:///./valera.db" - + + # XRay settings + xray_ssh_host: str = "mt.mirv.top" + xray_ssh_user: str = "root" + xray_ssh_password: str = "" + xray_add_user_script: str = "/root/bin/add_xray_user.sh" + class Config: env_file = ".env" extra = "allow" diff --git a/requirements.txt b/requirements.txt index 449a0d2..cc2c138 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,5 @@ vosk>=0.3.45 faster-whisper>=0.10.0 gigachat>=0.2.0 requests>=2.31.0 +asyncssh>=2.14.0 +qrcode[pil]>=7.4.0 diff --git a/src/bot/main.py b/src/bot/main.py index 035bf49..f4a8b10 100644 --- a/src/bot/main.py +++ b/src/bot/main.py @@ -3,23 +3,28 @@ import logging import os import re import subprocess -from telegram import Update, BotCommand +import io +from telegram import Update, BotCommand, InputFile from telegram.ext import ( - Application, CommandHandler, MessageHandler, filters, + Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler ) from telegram import InlineKeyboardButton, InlineKeyboardMarkup from config.config import get_settings from src.tools.orchestrator import Orchestrator -from src.bot.states import chat_state, ChatMode +from src.tools.xray import get_xray_client +from src.bot.states import chat_state, ChatMode, XRayState from src.bot.config_manager import get_selected_tool, get_selected_model, set_tool logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - level=logging.INFO + level=logging.WARNING ) logger = logging.getLogger(__name__) +# Отключаем подробное логирование asyncssh чтобы не засорять логи +logging.getLogger("asyncssh").setLevel(logging.ERROR) + settings = get_settings() orchestrator = Orchestrator() @@ -28,12 +33,48 @@ DANGEROUS_PATTERNS = [ r'\bсоздать\b', r'\bзаписать\b', r'\bудалить\b', r'\bизменить\b', r'\.write\(', r'\.save\(', r'\brm\b', r'\bmkdir\b', r'\bcp\b', r'\bсделать\b', r'\bвыполнить\b', r'\brun\b', r'\bзапустить\b', + r'\bпоказать\b', r'\bпосмотреть\b', r'\bпокажи\b', r'\bглянь\b', r'sudo', r'pip install', r'apt install', r'yum install', + r'chmod', r'chown', r'systemctl', r'service ', + r'curl.*\|', r'wget.*\|', r'bash.*-c', + r'exec', r'eval', r'shell', r'terminal', + r'\bgit\b', r'\bгит\b', r'\bnpm\b', r'\byarn\b', r'\bgo run\b', r'\bpip\b', + r'\bstatus\b', r'\bstat\b', r'\bстатус\b', ] +COMMAND_PATTERNS = [ + r'^ls\b', r'^dir\b', r'^cd\b', r'^pwd\b', + r'^cat\b', r'^grep\b', r'^find\b', r'^touch\b', + r'^mkdir\b', r'^rm\b', r'^cp\b', r'^mv\b', + r'^chmod\b', r'^chown\b', + r'^git\b', + r'list files', r'show files', r'покажи файлы', + r'выполни команду', r'run command', r'запусти команду', +] + +def is_simple_command_prompt(prompt: str) -> bool: + """Check if prompt is a simple system command that should be auto-executed""" + prompt_stripped = prompt.strip().lower() + simple_commands = [ + 'ls', 'dir', 'pwd', 'whoami', 'date', 'uname', + 'git status', 'git log', 'git diff', 'git branch', + 'df -h', 'free -h', 'top', 'ps aux' + ] + + for cmd in simple_commands: + if prompt_stripped == cmd or prompt_stripped.startswith(cmd + ' '): + return True + + # Check against patterns + for pattern in COMMAND_PATTERNS: + if re.search(pattern, prompt_stripped, re.IGNORECASE): + return True + + return False + def is_dangerous(prompt: str) -> bool: prompt_lower = prompt.lower() - for pattern in DANGEROUS_PATTERNS: + for pattern in DANGEROUS_PATTERNS + COMMAND_PATTERNS: if re.search(pattern, prompt_lower, re.IGNORECASE): return True return False @@ -42,7 +83,7 @@ def is_dangerous(prompt: str) -> bool: async def get_opencode_models(): try: result = subprocess.run( - ["opencode", "models"], + ["/home/mirivlad/.opencode/bin/opencode", "models"], capture_output=True, text=True, timeout=30 @@ -82,7 +123,8 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE): "/qwen - Использовать Qwen (все запросы через Qwen)\n" "/open - Выбрать модель OpenCode\n" "/mode confirm/auto - Режим подтверждения\n" - "/forget - Очистить историю\n\n" + "/forget - Очистить историю\n" + "/xray [email] - Добавить пользователя XRay\n\n" f"🔧 Текущая модель: {current_tool}" ) if model: @@ -116,32 +158,57 @@ async def confirm_callback(update: Update, context: ContextTypes.DEFAULT_TYPE): data = query.data if data.startswith("model_"): - model = data.replace("model_", "") + model_name = data.replace("model_", "") + if not model_name.startswith("opencode/"): + model = "opencode/" + model_name + else: + model = model_name set_tool("opencode", model) - await query.edit_message_text(f"✅ Выбрана модель OpenCode: {model}\nВсе запросы будут идти через эту модель.") + await query.edit_message_text(f"✅ Выбрана модель: {model}\nВсе запросы будут идти через эту модель.") return if data == "confirm_yes": - await query.edit_message_text("✅ Подтверждено. Выполняю...") pending = chat_state.get_pending_action(chat_id) - if pending: - prompt = pending.get("prompt") - tool = pending.get("tool") - chat_state.set_waiting_confirmation(chat_id, False) - await execute_tool_query(update, tool, prompt) + if not pending: + await query.edit_message_text("❌ Ошибка: действие не найдено") + return + + await query.edit_message_text("✅ Подтверждено. Выполняю...") + prompt = pending.get("prompt") + tool = pending.get("tool") + dangerous = pending.get("dangerous", False) + chat_state.set_waiting_confirmation(chat_id, False) + + model_raw = get_selected_model() + model = None + if tool == "opencode" and model_raw: + if model_raw.startswith("opencode/"): + model = model_raw.replace("opencode/", "", 1) + else: + model = model_raw + + thinking_msg = await query.message.reply_text("🤔 Думаю...") + + simple_prompt = prompt + result, success = await orchestrator.ask(simple_prompt, chat_id, tool, model, yolo=True) + + text = result[:4096] if len(result) > 4096 else result + await query.message.reply_text(text, parse_mode="Markdown") + try: + await thinking_msg.delete() + except: + pass elif data == "confirm_no": chat_state.set_waiting_confirmation(chat_id, False) - await query.edit_message_text("❌ Отменено.") + await query.edit_message_text("❌ Отменено. Команда не будет выполнена.") 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 - await update.message.reply_text("🤔 Думаю...") - model = get_selected_model() if tool == "opencode" else None - result, success = await orchestrator.ask(prompt, chat_id, tool, model) + result, success = await orchestrator.ask(prompt, chat_id, tool, model, yolo=False) text = result[:4096] if len(result) > 4096 else result @@ -163,7 +230,18 @@ async def qwen_command(update: Update, context: ContextTypes.DEFAULT_TYPE): return prompt = " ".join(context.args) - await execute_tool_query(update, "qwen", prompt) + chat_id = update.effective_chat.id + mode = chat_state.get_mode(chat_id) + use_yolo = mode == ChatMode.AUTO + + thinking_msg = await update.message.reply_text("🤔 Думаю..." + (" (YOLO)" if use_yolo else "")) + result, success = await orchestrator.ask(prompt, chat_id, "qwen", None, yolo=use_yolo) + text = result[:4096] if len(result) > 4096 else result + await update.message.reply_text(text) + try: + await thinking_msg.delete() + except: + pass async def open_menu(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -193,13 +271,24 @@ async def open_command(update: Update, context: ContextTypes.DEFAULT_TYPE): model = get_selected_model() if tool != "opencode": - set_tool("opencode", "minimax-m2.5-free") - model = "minimax-m2.5-free" + set_tool("opencode", "opencode/minimax-m2.5-free") + model = "opencode/minimax-m2.5-free" if not model: - model = "minimax-m2.5-free" + model = "opencode/minimax-m2.5-free" - await execute_tool_query(update, "opencode", prompt) + chat_id = update.effective_chat.id + mode = chat_state.get_mode(chat_id) + use_yolo = mode == ChatMode.AUTO + + thinking_msg = await update.message.reply_text("🤔 Думаю..." + (" (YOLO)" if use_yolo else "")) + result, success = await orchestrator.ask(prompt, chat_id, "opencode", model, yolo=use_yolo) + text = result[:4096] if len(result) > 4096 else result + await update.message.reply_text(text) + try: + await thinking_msg.delete() + except: + pass async def forget_command(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -208,59 +297,158 @@ async def forget_command(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text("🗑️ История чата очищена.") +async def xray_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + """ + Команда /xray - добавление пользователя XRay + 0. Ожидает ввод email + 1. Подключается по SSH к серверу + 2. Запускает скрипт add_xray_user.sh + 3. Возвращает результат и QR-код + """ + chat_id = update.effective_chat.id + + # Проверяем есть ли аргументы команды + if context.args: + # Если email передан сразу командой /xray email@example.com + email = context.args[0] + await process_xray_email(update, context, email) + return + + # Иначе запрашиваем email + chat_state.set_xray_state(chat_id, XRayState.WAITING_EMAIL) + await update.message.reply_text( + "🔐 *XRay - добавление пользователя*\n\n" + "Введите email для нового пользователя:" + ) + + +async def process_xray_email(update: Update, context: ContextTypes.DEFAULT_TYPE, email: str): + """Обработка email и выполнение скрипта на сервере""" + chat_id = update.effective_chat.id + + # Простая валидация email + email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$' + if not re.match(email_pattern, email): + await update.message.reply_text( + "❌ Неверный формат email. Попробуйте еще раз:" + ) + return + + # Очищаем состояние + chat_state.set_xray_state(chat_id, None) + + progress_msg = await update.message.reply_text( + f"⏳ Подключаюсь к серверу и добавляю пользователя {email}..." + ) + + try: + xray_client = get_xray_client() + + # Подключение + connected = await xray_client.connect() + if not connected: + await progress_msg.edit_text( + "❌ Ошибка подключения к серверу. Проверьте настройки SSH." + ) + return + + # Добавление пользователя + success, output, qr_png = await xray_client.add_user(email) + + # Отключение + await xray_client.disconnect() + + if not success: + await progress_msg.edit_text(f"❌ Ошибка:\n{output}") + return + + # Извлекаем vless ссылку из вывода + vless_pattern = r'(vless://[^\s]+)' + vless_match = re.search(vless_pattern, output) + vless_link = vless_match.group(1) if vless_match else None + + # Отправляем результат + await progress_msg.delete() + + result_text = ( + f"Привет! Вот данные для добавления VPN:\n\n" + f"1️⃣ Скачать приложение Hiddify тут:\n" + f"https://github.com/hiddify/hiddify-app/releases\n" + f"или в PlayMarket\n\n" + f"2️⃣ Добавить профиль в программе по ссылке:\n" + f"`{vless_link}`\n\n" + f"либо отсканируй QR-код прямо в программе" + ) + + await update.message.reply_text(result_text, parse_mode="Markdown") + + # Отправляем QR-код если есть + if qr_png: + await update.message.reply_photo( + photo=InputFile(io.BytesIO(qr_png), filename="xray_qr.png"), + caption=f"QR-код для {email}" + ) + + except Exception as e: + logger.error(f"Ошибка при добавлении пользователя XRay: {e}") + await progress_msg.edit_text(f"❌ Произошла ошибка: {e}") + + +async def handle_xray_email(update: Update, context: ContextTypes.DEFAULT_TYPE): + """Обработчик ввода email для команды /xray""" + chat_id = update.effective_chat.id + state = chat_state.get_xray_state(chat_id) + + if state != XRayState.WAITING_EMAIL: + return # Не в состоянии ожидания email + + email = update.message.text.strip() + await process_xray_email(update, context, email) + + 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) tool = get_selected_tool() - - dangerous = is_dangerous(prompt) - - if mode == ChatMode.CONFIRM and dangerous: - 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": tool, - "prompt": prompt - }) - - await update.message.reply_text( - f"⚠️ Это действие может внести изменения. Выполнить?\n\n{prompt[:200]}...", - reply_markup=reply_markup - ) - else: - await update.message.reply_text("🤔 Думаю...") - model = get_selected_model() if tool == "opencode" else None - await execute_tool_query(update, tool, prompt) + model = get_selected_model() if tool == "opencode" else None + + # Сначала проверяем не ждем ли мы email для XRay + if chat_state.get_xray_state(chat_id) == XRayState.WAITING_EMAIL: + await handle_xray_email(update, context) + return + + thinking_msg = await update.message.reply_text("🤔 Думаю...") + + result, success = await orchestrator.ask(prompt, chat_id, tool, model, yolo=True) + + text = result.strip() if result else "" + if not text: + text = "⚠️ Пустой ответ от модели." + + text = text[:4096] + await thinking_msg.edit_text(text, parse_mode="Markdown") def main(): builder = Application.builder() builder.token(settings.telegram_bot_token) - + if settings.telegram_proxy_url: builder = builder.proxy(settings.telegram_proxy_url) - logger.info(f"Используется прокси: {settings.telegram_proxy_url}") - + application = builder.build() - + application.add_handler(CommandHandler("start", start)) application.add_handler(CommandHandler("help", help_command)) application.add_handler(CommandHandler("mode", mode_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(CommandHandler("xray", xray_command)) application.add_handler(CallbackQueryHandler(confirm_callback)) application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message)) - - logger.info("Бот запущен") + application.run_polling(allowed_updates=Update.ALL_TYPES) diff --git a/src/bot/states.py b/src/bot/states.py index 88150ac..057a93e 100644 --- a/src/bot/states.py +++ b/src/bot/states.py @@ -10,22 +10,26 @@ class ChatMode(str, Enum): AUTO = "auto" +class XRayState(str, Enum): + """Состояния для команды /xray""" + WAITING_EMAIL = "waiting_email" + + 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] = {} @@ -34,18 +38,30 @@ class ChatState: 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") - + + # Методы для состояния XRay + def set_xray_state(self, chat_id: int, state: Optional[XRayState]): + if chat_id not in self.states: + self.states[chat_id] = {} + if state: + self.states[chat_id]["xray_state"] = state + else: + self.states[chat_id].pop("xray_state", None) + + def get_xray_state(self, chat_id: int) -> Optional[XRayState]: + return self.states.get(chat_id, {}).get("xray_state") + def clear(self, chat_id: int): self.states.pop(chat_id, None) diff --git a/src/memory/memory.py b/src/memory/memory.py index 2a19bfe..fc18740 100644 --- a/src/memory/memory.py +++ b/src/memory/memory.py @@ -25,14 +25,11 @@ class Memory: ) try: if os.path.exists(MODEL_PATH): - logger.info(f"Загрузка модели из локальной папки: {MODEL_PATH}") self.embedding_model = SentenceTransformer(MODEL_PATH) else: - logger.info("Загрузка модели с HuggingFace (локальная версия не найдена)") os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True) self.embedding_model = SentenceTransformer(MODEL_NAME) self.embedding_model.save(MODEL_PATH) - logger.info(f"Модель сохранена в: {MODEL_PATH}") except Exception as e: logger.warning(f"Не удалось загрузить модель эмбеддингов: {e}") self.embedding_model = None @@ -47,7 +44,6 @@ class Memory: ids=[doc_id], metadatas=[{"chat_id": str(chat_id), "role": role}] ) - logger.info(f"Добавлено сообщение в чат {chat_id}: {role}") def get_recent_messages(self, chat_id: int, limit: int = None) -> List[Dict]: if limit is None: @@ -94,12 +90,11 @@ class Memory: def clear_chat(self, chat_id: int): chat_id_str = self._get_chat_id(chat_id) - + results = self.collection.get(where={"chat_id": chat_id_str}) - + if results and results.get("ids"): self.collection.delete(ids=results["ids"]) - logger.info(f"Очищена история чата {chat_id}") def get_context_for_prompt(self, chat_id: int) -> str: messages = self.get_recent_messages(chat_id) diff --git a/src/scheduler/scheduler.py b/src/scheduler/scheduler.py index 898df01..29b7bc9 100644 --- a/src/scheduler/scheduler.py +++ b/src/scheduler/scheduler.py @@ -18,8 +18,6 @@ class SchedulerManager: self.reminders: Dict[int, List[dict]] = {} async def generate_idea(self): - logger.info("Запуск генерации идей") - for chat_id in self.orchestrator.memory.collection.get().get("metadatas", []): try: chat_id_int = int(chat_id.get("chat_id", 0)) @@ -57,30 +55,26 @@ class SchedulerManager: replace_existing=True ) self.scheduler.start() - logger.info("Планировщик запущен") - + def stop(self): self.scheduler.shutdown() - logger.info("Планировщик остановлен") def add_reminder(self, chat_id: int, text: str, run_at: datetime): if chat_id not in self.reminders: self.reminders[chat_id] = [] - + job = self.scheduler.add_job( self._send_reminder, trigger=DateTrigger(run_date=run_at), args=[chat_id, text], id=f"reminder_{chat_id}_{len(self.reminders[chat_id])}" ) - + self.reminders[chat_id].append({ "job_id": job.id, "text": text, "run_at": run_at }) - - logger.info(f"Напоминание добавлено для чата {chat_id}: {text} в {run_at}") async def _send_reminder(self, chat_id: int, text: str): try: diff --git a/src/speech/speech.py b/src/speech/speech.py index 71401b2..e4f8749 100644 --- a/src/speech/speech.py +++ b/src/speech/speech.py @@ -19,28 +19,25 @@ class SpeechRecognizer: def load_model(self): if not self.enabled: - logger.info("Распознавание речи отключено") return - + try: if self.model_name == "vosk": from vosk import Model, KaldiRecognizer import json - + model_path = os.path.expanduser("~/.vosk/models/vosk-model-ru") if not os.path.exists(model_path): logger.warning(f"Модель Vosk не найдена по пути {model_path}") return - + self.model = Model(model_path) - logger.info("Модель Vosk загружена") - + elif self.model_name == "whisper": from faster_whisper import WhisperModel - + self.model = WhisperModel("small", device="cpu", compute_type="int8") - logger.info("Модель Whisper загружена") - + except Exception as e: logger.error(f"Ошибка загрузки модели распознавания: {e}") self.enabled = False @@ -112,7 +109,6 @@ class SpeechRecognizer: def toggle(self, enabled: bool): self.enabled = enabled - logger.info(f"Распознавание речи: {'включено' if enabled else 'отключено'}") - + def is_enabled(self) -> bool: return self.enabled diff --git a/src/tools/orchestrator.py b/src/tools/orchestrator.py index 84b8991..9d8b446 100644 --- a/src/tools/orchestrator.py +++ b/src/tools/orchestrator.py @@ -8,9 +8,14 @@ from src.memory.memory import Memory logger = logging.getLogger(__name__) settings = get_settings() -SYSTEM_PROMPT = """Ты Валера - дружелюбный, умный и полезный программист-ассистент. -Ты помогаешь пользователям с программированием, отвечаешь на вопросы, объясняешь код и помогаешь решать задачи. -Будь кратким, но информативным. Используй кодовые блоки для примеров.""" +SYSTEM_PROMPT = """Ты Валера - дружелюбный Telegram-бот ассистент. + +ВАЖНО: Ты имеешь долгосрочную память! Ты помнишь ВСЕ предыдущие разговоры с этим пользователем из базы данных ChromaDB. +Когда пользователь спрашивает о твоей памяти - скажи что ты помнишь все ваши предыдущие разговоры. + +Ты помогаешь с программированием, отвечаешь на вопросы. +Будь кратким, используй кодовые блоки. +Твое имя - Валера.""" class Orchestrator: @@ -39,17 +44,30 @@ class Orchestrator: 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}" + similar = self.memory.search_similar(chat_id, user_prompt, limit=3) + if similar: + full_prompt += f"Похожие предыдущие разговоры:\n" + for sim in similar: + full_prompt += f"- {sim}\n" + full_prompt += "\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, model: Optional[str] = None) -> Tuple[str, bool]: + async def ask(self, prompt: str, chat_id: int, tool: Optional[str] = None, model: Optional[str] = None, yolo: bool = False) -> Tuple[str, bool]: selected_tool = tool or self.default_tool if selected_tool == "opencode" and model: - selected_tool = f"opencode:{model}" + if not model.startswith("opencode/"): + model_id = f"opencode/{model}" + else: + model_id = model + selected_tool = f"opencode:{model_id}" full_prompt = self._build_prompt(prompt, chat_id) @@ -58,22 +76,31 @@ class Orchestrator: elif selected_tool == "yandex": result, success = await self.yandex.ask(full_prompt) else: - result, success = await self.tool_runner.run_tool(selected_tool, full_prompt, model) + if selected_tool == "qwen": + result, success = await self.tool_runner.run_qwen(full_prompt, yolo) + else: + result, success = await self.tool_runner.run_tool(selected_tool, full_prompt, model, yolo) + + tool_key = selected_tool.split(":")[0] 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 tool_key in self.tool_limits: + self.tool_limits[tool_key]["failed"] += 1 + + if self.tool_limits[tool_key]["failed"] >= self.tool_limits[tool_key]["max_failures"]: + alt_tool = "opencode" if selected_tool == "qwen" else "qwen" + + if alt_tool == "qwen": + result, success = await self.tool_runner.run_qwen(full_prompt, yolo) + else: + result, success = await self.tool_runner.run_tool(alt_tool, full_prompt, model, yolo) + selected_tool = alt_tool - if success: - self.tool_limits[selected_tool]["failed"] = 0 + tool_key = selected_tool.split(":")[0] + if success and tool_key in self.tool_limits: + self.tool_limits[tool_key]["failed"] = 0 self.memory.add_message(chat_id, "user", prompt) self.memory.add_message(chat_id, "assistant", result) @@ -83,8 +110,7 @@ class Orchestrator: def set_default_tool(self, tool: str): if tool in ["qwen", "opencode", "gigachat", "yandex"]: self.default_tool = tool - logger.info(f"Инструмент по умолчанию изменён на {tool}") - + def get_default_tool(self) -> str: return self.default_tool diff --git a/src/tools/tool_runner.py b/src/tools/tool_runner.py index f731549..9a5f3a1 100644 --- a/src/tools/tool_runner.py +++ b/src/tools/tool_runner.py @@ -13,25 +13,35 @@ class ToolRunner: self.opencode_command = settings.opencode_command self.timeout = settings.tool_timeout - async def run_qwen(self, prompt: str) -> Tuple[str, bool]: - return await self._run_tool(self.qwen_command, prompt) + async def run_qwen(self, prompt: str, yolo: bool = False, cwd: str = None) -> Tuple[str, bool]: + import os + if cwd is None: + cwd = os.path.expanduser("~") + args = ["run", "--chat-recording=false", prompt] + if yolo: + args.append("--yolo") + return await self._run_tool(self.qwen_command, args, cwd=cwd) async def run_opencode(self, prompt: str, model: Optional[str] = None) -> Tuple[str, bool]: cmd = self.opencode_command if model: - cmd = f"{self.opencode_command}:{model}" - return await self._run_tool(cmd, prompt) + if not model.startswith("opencode/"): + model_id = f"opencode/{model}" + else: + model_id = model + args = ["run", "--model", model_id, prompt] + else: + args = ["run", prompt] + return await self._run_tool(cmd, args) - async def _run_tool(self, command: str, prompt: str) -> Tuple[str, bool]: - cmd_parts = command.split(":") - actual_cmd = cmd_parts[0] - + async def _run_tool(self, command: str, args: list, cwd: str = None) -> Tuple[str, bool]: try: process = await asyncio.create_subprocess_exec( - actual_cmd, - prompt, + command, + *args, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + stderr=asyncio.subprocess.PIPE, + cwd=cwd ) try: @@ -59,9 +69,9 @@ class ToolRunner: logger.exception("Ошибка при выполнении инструмента") return f"Ошибка: {str(e)}", False - async def run_tool(self, tool_name: str, prompt: str, model: Optional[str] = None) -> Tuple[str, bool]: + async def run_tool(self, tool_name: str, prompt: str, model: Optional[str] = None, yolo: bool = False) -> Tuple[str, bool]: if tool_name == "qwen": - return await self.run_qwen(prompt) + return await self.run_qwen(prompt, yolo) elif tool_name == "opencode": return await self.run_opencode(prompt, model) elif tool_name.startswith("opencode:"): diff --git a/src/tools/xray.py b/src/tools/xray.py new file mode 100644 index 0000000..d6ecf78 --- /dev/null +++ b/src/tools/xray.py @@ -0,0 +1,172 @@ +import asyncio +import re +import io +import base64 +from typing import Optional, Tuple +import logging + +import asyncssh +from PIL import Image +import qrcode + +logger = logging.getLogger(__name__) + + +class XRaySSHClient: + """Клиент для подключения к XRay серверу через SSH""" + + def __init__(self, host: str, user: str, password: str, script_path: str): + self.host = host + self.user = user + self.password = password + self.script_path = script_path + self.conn: Optional[asyncssh.SSHClientConnection] = None + + async def connect(self): + """Подключение к SSH серверу""" + try: + self.conn = await asyncssh.connect( + host=self.host, + username=self.user, + password=self.password, + known_hosts=None, # Отключаем проверку known_hosts для простоты + connect_timeout=15, + ) + return True + except asyncssh.PermissionDenied as e: + logger.error(f"SSH: доступ запрещён - {e}") + return False + except asyncssh.ConnectionLost as e: + logger.error(f"SSH: соединение потеряно - {e}") + return False + except Exception as e: + logger.error(f"Ошибка подключения к SSH: {e}") + return False + + async def disconnect(self): + """Отключение от SSH сервера""" + if self.conn: + self.conn.close() + await self.conn.wait_closed() + + async def add_user(self, email: str) -> Tuple[bool, str, Optional[bytes]]: + """ + Добавление пользователя через скрипт add_xray_user.sh + + Returns: + Tuple[success, output_text, qr_code_png] + """ + if not self.conn: + return False, "Нет подключения к серверу", None + + try: + # Скрипт читает email из stdin через read -p + # Поэтому передаем его через stdin + result = await self.conn.run( + self.script_path, + input=email + "\n", + timeout=120 + ) + + output = result.stdout + stderr = result.stderr + + full_output = output + if stderr: + full_output += "\n" + stderr + + if result.exit_status != 0: + return False, f"Ошибка скрипта (код {result.exit_status}):\n{full_output}", None + + # Парсим QR-код из вывода + qr_png = await self._extract_qr_code(full_output) + + return True, full_output, qr_png + + except asyncio.TimeoutError: + logger.error("Таймаут выполнения команды (120 сек)") + return False, "Таймаут выполнения команды (превышено 120 сек)", None + except Exception as e: + logger.error(f"Ошибка выполнения команды: {e}") + return False, f"Ошибка: {e}", None + + async def _extract_qr_code(self, output: str) -> Optional[bytes]: + """ + Извлечение QR-кода из вывода скрипта и конвертация в PNG + + Скрипт может выводить QR-код в разных форматах: + 1. ASCII art в консоли + 2. Base64 закодированное изображение + 3. Ссылка на подписку (vless://...) + + Возвращает PNG изображение или None + """ + # Пытаемся найти vless:// ссылку или другую subscription ссылку + vless_pattern = r'(vless://[^\s]+)' + trojan_pattern = r'(trojan://[^\s]+)' + vmess_pattern = r'(vmess://[^\s]+)' + + subscription_url = None + + for pattern in [vless_pattern, trojan_pattern, vmess_pattern]: + match = re.search(pattern, output) + if match: + subscription_url = match.group(1) + break + + if subscription_url: + # Генерируем QR-код из ссылки + try: + qr = qrcode.QRCode( + version=1, + error_correction=qrcode.constants.ERROR_CORRECT_L, + box_size=10, + border=2, + ) + qr.add_data(subscription_url) + qr.make(fit=True) + + img = qr.make_image(fill_color="black", back_color="white") + + # Конвертируем в bytes + img_bytes = io.BytesIO() + img.save(img_bytes, format='PNG') + img_bytes.seek(0) + + return img_bytes.getvalue() + + except Exception as e: + logger.error(f"Ошибка генерации QR-кода: {e}") + + # Если не нашли ссылку, пробуем найти base64 изображение + base64_pattern = r'(?:data:image/png;base64,)?([A-Za-z0-9+/]{100,}={0,2})' + base64_match = re.search(base64_pattern, output) + + if base64_match: + try: + base64_data = base64_match.group(1) + png_data = base64.b64decode(base64_data) + + # Проверяем что это валидное PNG + img = Image.open(io.BytesIO(png_data)) + img.verify() + + return png_data + + except Exception as e: + logger.error(f"Ошибка декодирования base64 QR-кода: {e}") + + return None + + +def get_xray_client() -> XRaySSHClient: + """Создание клиента из настроек окружения""" + from config.config import get_settings + settings = get_settings() + + return XRaySSHClient( + host=settings.xray_ssh_host, + user=settings.xray_ssh_user, + password=settings.xray_ssh_password, + script_path=settings.xray_add_user_script, + )