J.A.R.V.I.S-py/vision_handler.py

"""Vision handler — port of `crates/jarvis-core/src/lua/api/vision.rs`.

Captures the primary screen via PowerShell System.Drawing, base64-encodes
the PNG, sends to Groq's vision-capable model. Requires GROQ_TOKEN.

Action types in commands.yaml:
  vision_describe     "что на экране" / "опиши экран"
  vision_read_error   "прочитай ошибку"
"""

import base64
import os
import subprocess
import tempfile

_speak_fn = print


def set_speak(fn):
    global _speak_fn
    _speak_fn = fn


def _speak(text):
    try:
        _speak_fn(text)
    except Exception as exc:
        print(f"[vision] speak: {exc}")


_SCREENSHOT_PS_TEMPLATE = """
Add-Type -AssemblyName System.Windows.Forms;
Add-Type -AssemblyName System.Drawing;
$b = [System.Windows.Forms.Screen]::PrimaryScreen.Bounds;
$bmp = New-Object System.Drawing.Bitmap $b.Width, $b.Height;
$g = [System.Drawing.Graphics]::FromImage($bmp);
$g.CopyFromScreen($b.Location, [System.Drawing.Point]::Empty, $b.Size);
$bmp.Save('{path}', [System.Drawing.Imaging.ImageFormat]::Png);
$g.Dispose(); $bmp.Dispose();
"""


def _take_screenshot() -> str | None:
    """Save full screen to a temp PNG, return path or None on failure."""
    fd, path = tempfile.mkstemp(prefix='jarvis-screen-', suffix='.png')
    os.close(fd)
    ps = _SCREENSHOT_PS_TEMPLATE.format(path=path.replace("'", "''"))
    try:
        subprocess.run(
            ['powershell', '-NoProfile', '-ExecutionPolicy', 'Bypass', '-Command', ps],
            capture_output=True, timeout=10, check=True,
        )
    except Exception as exc:
        print(f"[vision] screenshot: {exc}")
        try:
            os.unlink(path)
        except OSError:
            pass
        return None
    if not os.path.isfile(path) or os.path.getsize(path) < 1000:
        try:
            os.unlink(path)
        except OSError:
            pass
        return None
    return path


def _vision_call(prompt: str, image_b64: str) -> str | None:
    # Vision is Groq-specific (Ollama doesn't expose vision via OpenAI-compat
    # endpoint in our stack). Use config.GROQ_TOKEN directly regardless of the
    # active text-LLM backend.
    try:
        import config as cfg
    except ImportError:
        return None
    if not getattr(cfg, 'GROQ_TOKEN', None):
        return None

    try:
        from openai import OpenAI
    except ImportError:
        return None

    client = OpenAI(
        api_key=cfg.GROQ_TOKEN,
        base_url=getattr(cfg, 'GROQ_BASE_URL', 'https://api.groq.com/openai/v1'),
    )

    model = os.environ.get('GROQ_VISION_MODEL', 'llama-3.2-11b-vision-preview')

    try:
        resp = client.chat.completions.create(
            model=model,
            messages=[{
                'role': 'user',
                'content': [
                    {'type': 'text', 'text': prompt},
                    {'type': 'image_url', 'image_url': {
                        'url': f'data:image/png;base64,{image_b64}',
                    }},
                ],
            }],
            max_tokens=400, temperature=0.2, timeout=30,
        )
        return resp.choices[0].message.content.strip()
    except Exception as exc:
        print(f"[vision] LLM call: {exc}")
        return None


def _describe(prompt: str) -> str | None:
    """Take a screenshot and ask vision-LLM. Returns description or None."""
    path = _take_screenshot()
    if not path:
        return None
    try:
        with open(path, 'rb') as f:
            b64 = base64.b64encode(f.read()).decode('ascii')
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass
    return _vision_call(prompt, b64)


def do_vision_describe(action, voice):
    _speak("Сейчас посмотрю.")
    prompt = (
        "Опиши коротко (1-3 предложения) что сейчас на экране. По-русски."
    )
    result = _describe(prompt)
    if not result:
        _speak("Не удалось разобрать экран. Проверь GROQ_TOKEN.")
        return
    _speak(result)


def do_vision_read_error(action, voice):
    _speak("Смотрю на ошибку.")
    prompt = (
        "Найди на экране текст ошибки (сообщение об ошибке, stack trace, диалог "
        "об исключении). Прочитай ключевое сообщение и одним коротким "
        "предложением подскажи возможную причину. Если ошибки нет — скажи "
        "'ошибок не вижу'. Отвечай по-русски. 1-3 предложения максимум."
    )
    result = _describe(prompt)
    if not result:
        _speak("Не получилось прочитать.")
        return
    _speak(result)