Roko
f71f3ce4be
- Two backends: RouterAI (gpt-4o-mini) + LocalGemma (Jan API) - Smart emoji routing (🌶️/🔞/🤫 → local Gemma, NSFW never hits cloud) - argparse with --help, --local, --model, --prompt - ImagePreprocessor with ffmpeg resize - ProxyHandler({}) for local server bypass - logging, try/finally cleanup, HTTPError/URLError handling - Zero shell injection, zero keys in /tmp - SKILL.md with full documentation for other agents
388 lines
15 KiB
Python
388 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""vision-analyze-v2.py — анализ изображений через RouterAI или локальную Gemma."""
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from abc import ABC, abstractmethod
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Константы
|
||
# ---------------------------------------------------------------------------
|
||
|
||
DEFAULT_PROMPT = (
|
||
"Дай последовательное описание композиции: передний план, задний план, "
|
||
"элементы на изображении, люди или персонажи. Пиши живым, человеческим языком, 3-5 предложений."
|
||
)
|
||
|
||
ROUTERAI_URL = "https://routerai.ru/api/v1/chat/completions"
|
||
ROUTERAI_DEFAULT_MODEL = "openai/gpt-4o-mini"
|
||
|
||
LOCAL_URL = "http://192.168.88.183:1337/messages"
|
||
LOCAL_MODEL = "gemma-4-E4B-it-Q8_0"
|
||
|
||
ROUTERAI_KEY_PATH = os.path.expanduser("~/.openclaw/openclaw.json")
|
||
LOCAL_KEY_PATH = os.path.expanduser("~/.openclaw/workspace/wireguard/jan_api_key.txt")
|
||
|
||
SPICY_EMOJIS = frozenset({"🌶️", "🔞", "🤫"})
|
||
ALLOWED_EXTENSIONS = frozenset({".jpg", ".jpeg", ".png", ".gif", ".webp"})
|
||
|
||
ROUTERAI_TIMEOUT = 30
|
||
LOCAL_TIMEOUT = 600
|
||
LOCAL_RETRY_DELAY = 5
|
||
|
||
RESIZE_THRESHOLD_BYTES = 1_000_000
|
||
RESIZE_SCALE = 1024
|
||
RESIZE_QUALITY = 3
|
||
|
||
LOGGER = logging.getLogger(__name__)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Помощники
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _load_routerai_key() -> str:
|
||
"""Читает API-ключ RouterAI из конфигурации openclaw."""
|
||
try:
|
||
with open(ROUTERAI_KEY_PATH, encoding="utf-8") as f:
|
||
config = json.load(f)
|
||
except FileNotFoundError as exc:
|
||
raise ValueError(f"Конфиг RouterAI не найден: {ROUTERAI_KEY_PATH}") from exc
|
||
except json.JSONDecodeError as exc:
|
||
raise ValueError(f"Невалидный JSON в конфиге RouterAI: {ROUTERAI_KEY_PATH}") from exc
|
||
|
||
key = config.get("models", {}).get("providers", {}).get("routerai-ru", {}).get("apiKey")
|
||
if not key:
|
||
raise ValueError("В конфиге RouterAI отсутствует models.providers['routerai-ru'].apiKey")
|
||
return key
|
||
|
||
|
||
def _load_local_key() -> str:
|
||
"""Читает API-ключ локального сервера Gemma."""
|
||
try:
|
||
with open(LOCAL_KEY_PATH, encoding="utf-8") as f:
|
||
key = f.read().strip()
|
||
except FileNotFoundError as exc:
|
||
raise ValueError(f"Ключ локального сервера не найден: {LOCAL_KEY_PATH}") from exc
|
||
|
||
if not key:
|
||
raise ValueError("Файл ключа локального сервера пуст")
|
||
return key
|
||
|
||
|
||
def _media_type_for_path(path: str) -> str:
|
||
"""Возвращает MIME-тип изображения по расширению."""
|
||
ext = os.path.splitext(path)[1].lower()
|
||
return {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".gif": "image/gif",
|
||
".webp": "image/webp",
|
||
}.get(ext, "image/jpeg")
|
||
|
||
|
||
def _encode_image(path: str, media_type: str | None = None) -> tuple[str, str]:
|
||
"""Кодирует изображение в base64 и возвращает (base64, media_type)."""
|
||
with open(path, "rb") as f:
|
||
data = f.read()
|
||
b64 = base64.b64encode(data).decode("ascii")
|
||
return b64, media_type or _media_type_for_path(path)
|
||
|
||
|
||
def _contains_spicy(prompt: str) -> bool:
|
||
"""Проверяет наличие NSFW-эмодзи в промпте."""
|
||
return any(emoji in prompt for emoji in SPICY_EMOJIS)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ImagePreprocessor
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class ImagePreprocessor:
|
||
"""Подготовка изображений перед отправкой в локальную модель."""
|
||
|
||
@staticmethod
|
||
def resize_if_needed(path: str) -> tuple[str, bool]:
|
||
"""
|
||
Если изображение больше порога, уменьшает его через ffmpeg.
|
||
|
||
Возвращает (effective_path, is_temporary). Если is_temporary=True,
|
||
вызывающий код должен удалить effective_path после использования.
|
||
"""
|
||
try:
|
||
size = os.path.getsize(path)
|
||
except OSError as exc:
|
||
raise ValueError(f"Не удалось получить размер файла: {path}") from exc
|
||
|
||
if size <= RESIZE_THRESHOLD_BYTES:
|
||
return path, False
|
||
|
||
tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
|
||
try:
|
||
tmp.close()
|
||
cmd = [
|
||
"ffmpeg",
|
||
"-y",
|
||
"-i", path,
|
||
"-vf", f"scale={RESIZE_SCALE}:-1",
|
||
"-q:v", str(RESIZE_QUALITY),
|
||
tmp.name,
|
||
]
|
||
proc = subprocess.run(cmd, capture_output=True, text=True)
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(f"ffmpeg завершился с ошибкой: {proc.stderr.strip()}")
|
||
if not os.path.isfile(tmp.name) or os.path.getsize(tmp.name) == 0:
|
||
raise RuntimeError("ffmpeg создал пустой файл")
|
||
return tmp.name, True
|
||
except Exception:
|
||
if os.path.exists(tmp.name):
|
||
os.unlink(tmp.name)
|
||
raise
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Бэкенды
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class Backend(ABC):
|
||
"""Абстрактный бэкенд для анализа изображений."""
|
||
|
||
@abstractmethod
|
||
def analyze(self, path: str, prompt: str) -> str:
|
||
"""Анализирует изображение по пути path с заданным prompt."""
|
||
...
|
||
|
||
|
||
class RouterAIBackend(Backend):
|
||
"""Бэкенд routerai.ru (OpenAI-совместимый API)."""
|
||
|
||
def __init__(self, model: str = ROUTERAI_DEFAULT_MODEL) -> None:
|
||
self.model = model
|
||
self._api_key = _load_routerai_key()
|
||
|
||
def analyze(self, path: str, prompt: str) -> str:
|
||
img_b64, media_type = _encode_image(path)
|
||
|
||
payload = json.dumps({
|
||
"model": self.model,
|
||
"modalities": ["image", "text"],
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {
|
||
"url": f"data:{media_type};base64,{img_b64}"
|
||
}},
|
||
],
|
||
}],
|
||
"max_tokens": 512,
|
||
}).encode("utf-8")
|
||
|
||
req = urllib.request.Request(
|
||
ROUTERAI_URL,
|
||
data=payload,
|
||
headers={
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
)
|
||
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=ROUTERAI_TIMEOUT) as resp:
|
||
data = json.loads(resp.read())
|
||
except urllib.error.HTTPError as exc:
|
||
body = exc.read().decode("utf-8", errors="replace")
|
||
LOGGER.error("RouterAI вернул HTTP %s: %s", exc.code, body)
|
||
raise RuntimeError(f"RouterAI HTTP {exc.code}: {body}") from exc
|
||
except urllib.error.URLError as exc:
|
||
LOGGER.error("Ошибка соединения с RouterAI: %s", exc.reason)
|
||
raise RuntimeError(f"Ошибка соединения с RouterAI: {exc.reason}") from exc
|
||
except json.JSONDecodeError as exc:
|
||
LOGGER.error("Невалидный JSON от RouterAI")
|
||
raise RuntimeError("Невалидный JSON от RouterAI") from exc
|
||
|
||
content = data.get("choices", [{}])[0].get("message", {}).get("content")
|
||
if content:
|
||
return content
|
||
|
||
LOGGER.error("Неожиданный ответ RouterAI: %s", data)
|
||
raise RuntimeError("В ответе RouterAI отсутствует content")
|
||
|
||
|
||
class LocalGemmaBackend(Backend):
|
||
"""Бэкенд локальной Gemma через Anthropic-совместимый endpoint."""
|
||
|
||
def __init__(self) -> None:
|
||
self._api_key = _load_local_key()
|
||
self._opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
|
||
|
||
@staticmethod
|
||
def _is_connection_error(exc: urllib.error.URLError) -> bool:
|
||
return isinstance(exc.reason, (ConnectionError, TimeoutError, socket.timeout))
|
||
|
||
def _build_request(self, payload: bytes) -> urllib.request.Request:
|
||
return urllib.request.Request(
|
||
LOCAL_URL,
|
||
data=payload,
|
||
headers={
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
"anthropic-version": "2023-06-01",
|
||
"Content-Type": "application/json",
|
||
},
|
||
)
|
||
|
||
def analyze(self, path: str, prompt: str) -> str:
|
||
effective_path, is_temporary = ImagePreprocessor.resize_if_needed(path)
|
||
|
||
try:
|
||
img_b64, _ = _encode_image(effective_path, media_type="image/jpeg")
|
||
|
||
payload = json.dumps({
|
||
"model": LOCAL_MODEL,
|
||
"max_tokens": 1024,
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image", "source": {
|
||
"type": "base64",
|
||
"media_type": "image/jpeg",
|
||
"data": img_b64,
|
||
}},
|
||
],
|
||
}],
|
||
}).encode("utf-8")
|
||
|
||
attempts = 0
|
||
while True:
|
||
req = self._build_request(payload)
|
||
try:
|
||
with self._opener.open(req, timeout=LOCAL_TIMEOUT) as resp:
|
||
data = json.loads(resp.read())
|
||
break
|
||
except urllib.error.HTTPError as exc:
|
||
body = exc.read().decode("utf-8", errors="replace")
|
||
LOGGER.error("Локальный сервер вернул HTTP %s: %s", exc.code, body)
|
||
raise RuntimeError(f"Локальный сервер HTTP {exc.code}: {body}") from exc
|
||
except urllib.error.URLError as exc:
|
||
attempts += 1
|
||
if attempts <= 1 and self._is_connection_error(exc):
|
||
LOGGER.warning(
|
||
"Соединение с локальным сервером сброшено, повторная попытка через %s сек...",
|
||
LOCAL_RETRY_DELAY,
|
||
)
|
||
time.sleep(LOCAL_RETRY_DELAY)
|
||
continue
|
||
LOGGER.error("Ошибка соединения с локальным сервером: %s", exc.reason)
|
||
raise RuntimeError(f"Ошибка соединения с локальным сервером: {exc.reason}") from exc
|
||
except json.JSONDecodeError as exc:
|
||
LOGGER.error("Невалидный JSON от локального сервера")
|
||
raise RuntimeError("Невалидный JSON от локального сервера") from exc
|
||
|
||
for block in data.get("content", []):
|
||
if block.get("type") == "text":
|
||
return block.get("text", "")
|
||
|
||
for block in data.get("content", []):
|
||
if block.get("type") == "thinking":
|
||
return block.get("thinking", "Только размышления, ответ не сформирован")
|
||
|
||
LOGGER.error("Неожиданный ответ локального сервера: %s", data)
|
||
raise RuntimeError("В ответе локального сервера отсутствуют текстовые блоки")
|
||
finally:
|
||
if is_temporary and os.path.exists(effective_path):
|
||
os.unlink(effective_path)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI и маршрутизация
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _resolve_backend(args: argparse.Namespace) -> Backend:
|
||
"""Выбирает бэкенд на основе аргументов и содержимого промпта."""
|
||
if args.local:
|
||
LOGGER.info("Форсирован локальный бэкенд Gemma")
|
||
return LocalGemmaBackend()
|
||
|
||
if args.model:
|
||
LOGGER.info("Используется RouterAI с моделью %s", args.model)
|
||
return RouterAIBackend(model=args.model)
|
||
|
||
if _contains_spicy(args.prompt):
|
||
LOGGER.info("Обнаружен spicy-флаг в промпте, используется локальный бэкенд Gemma")
|
||
return LocalGemmaBackend()
|
||
|
||
LOGGER.info("Используется RouterAI с моделью по умолчанию %s", ROUTERAI_DEFAULT_MODEL)
|
||
return RouterAIBackend()
|
||
|
||
|
||
def _validate_path(path: str) -> None:
|
||
"""Проверяет, что путь указывает на разрешённое изображение."""
|
||
if not os.path.isfile(path):
|
||
raise ValueError(f"Файл не найден: {path}")
|
||
|
||
ext = os.path.splitext(path)[1].lower()
|
||
if ext not in ALLOWED_EXTENSIONS:
|
||
raise ValueError(
|
||
f"Неподдерживаемый формат файла: {ext}. "
|
||
f"Разрешены: {', '.join(sorted(ALLOWED_EXTENSIONS))}"
|
||
)
|
||
|
||
|
||
def _build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
description="Анализ изображений через RouterAI или локальную Gemma."
|
||
)
|
||
parser.add_argument("path", help="Путь к изображению")
|
||
parser.add_argument("--local", action="store_true", help="Форсировать локальную Gemma")
|
||
parser.add_argument("--model", help="Модель для RouterAI (например, openai/gpt-4o-mini)")
|
||
parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Пользовательский промпт")
|
||
return parser
|
||
|
||
|
||
def main() -> int:
|
||
"""Точка входа скрипта."""
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S",
|
||
)
|
||
|
||
parser = _build_parser()
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
_validate_path(args.path)
|
||
except ValueError as exc:
|
||
LOGGER.error("%s", exc)
|
||
return 1
|
||
|
||
try:
|
||
backend = _resolve_backend(args)
|
||
result = backend.analyze(args.path, args.prompt)
|
||
except (RuntimeError, ValueError) as exc:
|
||
LOGGER.error("%s", exc)
|
||
return 1
|
||
|
||
print(result)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|