cosmo-voice-assistant/cosmo/tools.py

"""
Инструменты агента для smolagents.
Каждый инструмент — функция с декоратором @tool.
smolagents автоматически генерирует схему из docstring и type hints.
"""

import os
import subprocess
import webbrowser
import urllib.parse
from loguru import logger
from smolagents import tool

from cosmo.memory import Memory

# Глобальная ссылка на память — устанавливается из agent.py
_memory: Memory | None = None

def set_memory(mem: Memory):
    global _memory
    _memory = mem


# ------------------------------------------------------------------
# Поиск Git Bash
# ------------------------------------------------------------------

_GIT_BASH_CANDIDATES = [
    "C:/Program Files/Git/bin/bash.exe",
    "C:/Program Files (x86)/Git/bin/bash.exe",
    os.path.expandvars("%LOCALAPPDATA%/Programs/Git/bin/bash.exe"),
]

def _find_git_bash() -> str:
    for path in _GIT_BASH_CANDIDATES:
        if os.path.exists(path):
            return path
    try:
        result = subprocess.run(["where", "bash"], capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            first = result.stdout.strip().splitlines()[0]
            if "git" in first.lower():
                return first
    except Exception:
        pass
    return "bash"


# ------------------------------------------------------------------
# Инструменты
# ------------------------------------------------------------------

@tool
def run_shell(command: str) -> str:
    """
    Выполнить команду в Git Bash и получить вывод.
    Используй bash-синтаксис. Для запуска Windows-программ: start '' '/c/path/app.exe'
    или cmd //c start. Для поиска программ: which, find, cmd //c where.

    Args:
        command: bash команда для выполнения
    """
    bash = _find_git_bash()
    logger.info(f"[run_shell] {command}")
    try:
        result = subprocess.run(
            [bash, "-c", command],
            capture_output=True,
            text=True,
            timeout=20,
            encoding="utf-8",
            errors="replace",
        )
        output = result.stdout.strip()
        stderr = result.stderr.strip()
        if result.returncode != 0 and stderr:
            return f"[returncode={result.returncode}]\nstdout: {output}\nstderr: {stderr}"
        return output if output else f"[выполнено, returncode={result.returncode}]"
    except FileNotFoundError:
        return "[ошибка: Git Bash не найден]"
    except subprocess.TimeoutExpired:
        return "[таймаут 20с]"
    except Exception as e:
        return f"[ошибка: {e}]"


@tool
def find_program(name: str) -> str:
    """
    Найти программу или файл на Windows по имени.
    Ищет в PATH, Program Files, AppData и реестре.
    Используй когда не знаешь точный путь к программе.

    Args:
        name: имя программы без расширения, например 'webstorm' или 'chrome'
    """
    stem = name.lower().strip().replace(".exe", "")
    logger.info(f"[find_program] {stem}")

    steps = [
        f"which {stem} 2>/dev/null",
        f"cmd //c where {stem} 2>/dev/null",
        f"find '/c/Program Files' -iname '{stem}*.exe' 2>/dev/null | head -3",
        f"find '/c/Program Files (x86)' -iname '{stem}*.exe' 2>/dev/null | head -3",
        f"find \"$LOCALAPPDATA\" -iname '{stem}*.exe' 2>/dev/null | head -3",
        f"cmd //c 'reg query \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\App Paths\\{stem}.exe\" /ve 2>nul'",
    ]

    results = []
    for cmd in steps:
        out = run_shell(cmd)
        if out and not out.startswith("[") and len(out) > 3:
            results.append(out.strip())

    if results:
        return "Найдено:\n" + "\n".join(results)
    return f"Программа '{name}' не найдена. Попробуй другое имя."


@tool
def open_browser(url: str, search: bool = False) -> str:
    """
    Открыть URL в браузере или выполнить поиск в Google.

    Args:
        url: полный URL (https://...) или поисковый запрос если search=True
        search: если True — выполнить поиск Google по тексту в url
    """
    if search:
        url = "https://www.google.com/search?q=" + urllib.parse.quote(url)
    elif not url.startswith(("http://", "https://")):
        url = "https://" + url
    webbrowser.open(url)
    logger.info(f"[open_browser] {url}")
    return f"Открыт браузер: {url}"


@tool
def read_file(path: str) -> str:
    """
    Прочитать содержимое текстового файла.

    Args:
        path: абсолютный путь к файлу
    """
    try:
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            return f.read(8000)
    except FileNotFoundError:
        return f"Файл не найден: {path}"
    except Exception as e:
        return f"Ошибка чтения: {e}"


@tool
def write_file(path: str, content: str) -> str:
    """
    Записать текст в файл (создаёт или перезаписывает).

    Args:
        path: абсолютный путь к файлу
        content: содержимое файла
    """
    try:
        dir_path = os.path.dirname(path)
        if dir_path:
            os.makedirs(dir_path, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"Файл записан: {path}"
    except Exception as e:
        return f"Ошибка записи: {e}"


@tool
def memory_set(key: str, value: str) -> str:
    """
    Сохранить факт в долгосрочную память ассистента.
    Используй для запоминания путей программ, предпочтений пользователя, часто используемых команд.
    Ключи: 'app.название', 'user.имя', 'pref.что-то'.

    Args:
        key: ключ факта, например 'app.webstorm' или 'user.name'
        value: значение для сохранения
    """
    if _memory is None:
        return "Память не инициализирована"
    _memory.set(key, value)
    return f"Запомнено: {key} = {value!r}"


@tool
def memory_get(key: str) -> str:
    """
    Получить сохранённый факт из долгосрочной памяти по ключу.

    Args:
        key: ключ факта, например 'app.webstorm'
    """
    if _memory is None:
        return "Память не инициализирована"
    value = _memory.get(key)
    return f"{key} = {value!r}" if value else f"Факт '{key}' не найден"


@tool
def memory_list(prefix: str = "") -> str:
    """
    Показать все факты из памяти, опционально по префиксу.

    Args:
        prefix: префикс для фильтрации, например 'app.' чтобы видеть пути программ
    """
    if _memory is None:
        return "Память не инициализирована"
    facts = _memory.list_facts(prefix)
    if not facts:
        return "Память пуста" if not prefix else f"Нет фактов с префиксом '{prefix}'"
    return "\n".join(f"{k}: {v}" for k, v in sorted(facts.items()))


# Список всех инструментов для передачи в агент
ALL_TOOLS = [
    run_shell,
    find_program,
    open_browser,
    read_file,
    write_file,
    memory_set,
    memory_get,
    memory_list,
]