Mavis
2394eff1c0
- LM Studio client (httpx-based, OpenAI-compatible) - SVG validator (lxml, whitelist tags, no <script>/<foreignObject>/http refs) - PNG renderer (resvg-py primary, cairosvg fallback - no native cairo dep) - History (SQLite, tracks raw/validated/preview paths) - Gradio UI on 127.0.0.1:8788 with: * mode radio (icon/illustration) * n_candidates slider (default 1) * image upload for image-to-SVG * LM Studio URL/token inputs * model dropdown + refresh button - prompts/ with system_icon.txt, system_illustration.txt, few_shot_examples.txt - docs/spec.md, docs/design.md - 122 unit/integration tests passing
347 lines
14 KiB
Python
347 lines
14 KiB
Python
"""Клиент к LM Studio через OpenAI-compatible chat completions.
|
||
|
||
Использует синхронный httpx, чтобы не тащить openai SDK и не ловить
|
||
несовместимости их версий. Поддерживает текстовые и мультимодальные
|
||
(text + image_url) сообщения, а также n>1 для batch-генерации кандидатов.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import io
|
||
import logging
|
||
import os
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Iterable
|
||
|
||
import httpx
|
||
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
|
||
DEFAULT_BASE_URL = "http://127.0.0.1:1234/v1"
|
||
DEFAULT_MODEL = "qwen/qwen3.5-35b-a3b"
|
||
DEFAULT_API_KEY = "lm-studio" # LM Studio игнорирует значение, но заголовок обязателен
|
||
DEFAULT_TIMEOUT_S = 300.0 # qwen3.5 может долго рассуждать, не режем по таймауту
|
||
DEFAULT_MAX_TOKENS = 32768 # без жёсткого лимита: qwen3.5 ест reasoning + content в одном budget
|
||
|
||
|
||
class LMStudioUnavailable(RuntimeError):
|
||
"""LM Studio не отвечает, вернул 5xx или вернул неожиданный формат."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class LMTurnResult:
|
||
"""Результат одного запроса к LM Studio.
|
||
|
||
Attributes:
|
||
raw_texts: список текстов ответов — по одному на кандидат. Если модель
|
||
вернула tool_calls без `content`, элемент будет пустой строкой.
|
||
elapsed_s: время запроса в секундах.
|
||
model: фактическое имя модели, которое вернул сервер (если есть) или
|
||
то, что мы послали.
|
||
usage: словарь usage от сервера, может быть None.
|
||
finish_reasons: список finish_reason по кандидатам ("stop"/"length"/...).
|
||
"""
|
||
|
||
raw_texts: list[str]
|
||
elapsed_s: float
|
||
model: str
|
||
usage: dict | None = None
|
||
finish_reasons: list[str] = field(default_factory=list)
|
||
|
||
|
||
def _coerce_text_part(part: Any) -> str:
|
||
"""Достаёт текст из элемента content — поддерживает str и list[dict]."""
|
||
if isinstance(part, str):
|
||
return part
|
||
if isinstance(part, dict):
|
||
return str(part.get("text", "") or "")
|
||
return str(part)
|
||
|
||
|
||
def _normalize_messages(messages: Iterable[dict]) -> list[dict]:
|
||
"""Готовит сообщения к отправке: текст либо строкой, либо [{type:text},...].
|
||
|
||
Допускаем на входе content как str или list[dict] (text + image_url).
|
||
Возвращаем список в формате OpenAI chat completions.
|
||
"""
|
||
out: list[dict] = []
|
||
for msg in messages:
|
||
role = msg.get("role")
|
||
content = msg.get("content")
|
||
if role is None or content is None:
|
||
log.warning("пропускаю сообщение без role/content: %r", msg)
|
||
continue
|
||
if isinstance(content, str):
|
||
out.append({"role": role, "content": content})
|
||
else:
|
||
parts: list[dict] = []
|
||
for piece in content:
|
||
if isinstance(piece, str):
|
||
parts.append({"type": "text", "text": piece})
|
||
elif isinstance(piece, dict):
|
||
ptype = piece.get("type")
|
||
if ptype == "text":
|
||
parts.append({"type": "text", "text": str(piece.get("text", ""))})
|
||
elif ptype == "image_url":
|
||
url = piece.get("image_url")
|
||
if isinstance(url, dict):
|
||
url = url.get("url")
|
||
parts.append({"type": "image_url", "image_url": {"url": str(url)}})
|
||
else:
|
||
log.warning("неизвестный тип части content: %r", ptype)
|
||
out.append({"role": role, "content": parts})
|
||
return out
|
||
|
||
|
||
def chat(
|
||
*,
|
||
messages: list[dict],
|
||
model: str | None = None,
|
||
n: int = 1,
|
||
temperature: float = 0.4,
|
||
max_tokens: int = DEFAULT_MAX_TOKENS,
|
||
base_url: str | None = None,
|
||
api_key: str | None = None,
|
||
timeout_s: float | None = None,
|
||
) -> LMTurnResult:
|
||
"""Шлёт chat completion в LM Studio и возвращает N текстов ответов.
|
||
|
||
Args:
|
||
messages: список сообщений в формате OpenAI. content может быть str
|
||
или list[dict] (text/image_url).
|
||
model: имя модели; если None — берём DEFAULT_MODEL / env `DEFAULT_MODEL`.
|
||
n: количество кандидатов, 1..8.
|
||
temperature: 0.0..1.5.
|
||
max_tokens: верхняя граница длины ответа.
|
||
base_url: эндпоинт LM Studio; дефолт http://127.0.0.1:1234/v1.
|
||
api_key: Bearer-токен; дефолт "lm-studio".
|
||
timeout_s: общий таймаут httpx, дефолт 120 секунд.
|
||
|
||
Returns:
|
||
LMTurnResult с N текстами и метаданными.
|
||
|
||
Raises:
|
||
LMStudioUnavailable: при сетевых/HTTP/ошибках парсинга.
|
||
"""
|
||
if n < 1:
|
||
raise ValueError(f"n должно быть >= 1, получено {n}")
|
||
|
||
base = (base_url or os.environ.get("LM_STUDIO_BASE_URL") or DEFAULT_BASE_URL).rstrip("/")
|
||
key = api_key if api_key is not None else os.environ.get("LM_STUDIO_API_KEY", DEFAULT_API_KEY)
|
||
mdl = model or os.environ.get("DEFAULT_MODEL", DEFAULT_MODEL)
|
||
timeout = float(
|
||
os.environ.get("REQUEST_TIMEOUT_S", str(timeout_s if timeout_s is not None else DEFAULT_TIMEOUT_S))
|
||
)
|
||
|
||
url = f"{base}/chat/completions"
|
||
payload: dict[str, Any] = {
|
||
"model": mdl,
|
||
"messages": _normalize_messages(messages),
|
||
"n": n,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": False,
|
||
# попытка отключить thinking у qwen3.5; если LM Studio/модель не уважают — не страшно
|
||
"chat_template_kwargs": {"enable_thinking": False},
|
||
}
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {key}",
|
||
}
|
||
|
||
log.info(
|
||
"LM Studio → %s model=%s n=%d temp=%.2f timeout=%.0fs",
|
||
url, mdl, n, temperature, timeout,
|
||
)
|
||
started = time.monotonic()
|
||
try:
|
||
with httpx.Client(timeout=timeout) as client:
|
||
resp = client.post(url, json=payload, headers=headers)
|
||
except httpx.TimeoutException as exc:
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio timeout at {url}: запрос превысил {timeout:.0f}с"
|
||
) from exc
|
||
except httpx.HTTPError as exc:
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio недоступен по адресу {url}: {exc}"
|
||
) from exc
|
||
|
||
elapsed = time.monotonic() - started
|
||
|
||
if resp.status_code >= 500:
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio error: {resp.status_code} {resp.text[:200]}"
|
||
)
|
||
if resp.status_code >= 400:
|
||
# 4xx — это наша ошибка (плохой запрос / модель не загружена / квота).
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio вернул {resp.status_code}: {resp.text[:200]}"
|
||
)
|
||
|
||
try:
|
||
data = resp.json()
|
||
except ValueError as exc:
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio вернул не-JSON: {resp.text[:200]}"
|
||
) from exc
|
||
|
||
choices = data.get("choices") or []
|
||
if not choices:
|
||
raise LMStudioUnavailable(
|
||
f"LM Studio не вернул ни одного choice: {resp.text[:200]}"
|
||
)
|
||
|
||
raw_texts: list[str] = []
|
||
finish_reasons: list[str] = []
|
||
for ch in choices:
|
||
msg = ch.get("message") or {}
|
||
content = msg.get("content")
|
||
if isinstance(content, str):
|
||
raw_texts.append(content)
|
||
elif isinstance(content, list):
|
||
raw_texts.append("".join(_coerce_text_part(p) for p in content))
|
||
elif content is None:
|
||
# Модель могла вернуть tool_calls — для нас это пустой кандидат.
|
||
raw_texts.append("")
|
||
else:
|
||
raw_texts.append(_coerce_text_part(content))
|
||
finish_reasons.append(str(ch.get("finish_reason") or ""))
|
||
|
||
# Если сервер вернул меньше choice'ов, чем n — добиваем пустыми строками.
|
||
while len(raw_texts) < n:
|
||
raw_texts.append("")
|
||
finish_reasons.append("missing")
|
||
|
||
return LMTurnResult(
|
||
raw_texts=raw_texts,
|
||
elapsed_s=elapsed,
|
||
model=str(data.get("model") or mdl),
|
||
usage=data.get("usage"),
|
||
finish_reasons=finish_reasons,
|
||
)
|
||
|
||
|
||
def encode_pil_to_data_url(image: Any, *, mime: str = "image/png") -> str:
|
||
"""Кодирует PIL-картинку в data: URL для передачи в image_url.
|
||
|
||
Args:
|
||
image: объект PIL.Image.
|
||
mime: MIME-тип, по умолчанию image/png (LM Studio VLM лучше работают с PNG).
|
||
|
||
Returns:
|
||
Строка вида `data:image/png;base64,<...>`.
|
||
"""
|
||
from PIL import Image # локальный импорт, чтобы не требовать Pillow без надобности
|
||
|
||
buf = io.BytesIO()
|
||
img = image
|
||
if getattr(img, "mode", None) == "RGBA" and mime == "image/jpeg":
|
||
img = img.convert("RGB")
|
||
elif getattr(img, "mode", None) not in ("RGB", "RGBA", "L") and mime == "image/png":
|
||
img = img.convert("RGBA")
|
||
img.save(buf, format="PNG" if mime == "image/png" else "JPEG", optimize=True)
|
||
encoded = base64.b64encode(buf.getvalue()).decode("ascii")
|
||
return f"data:{mime};base64,{encoded}"
|
||
|
||
|
||
def validate_image(image: Any, *, max_bytes: int = 10 * 1024 * 1024, max_side: int = 4096) -> None:
|
||
"""Проверяет картинку перед отправкой в LM Studio (US-2, edge case #3).
|
||
|
||
Args:
|
||
image: PIL.Image.
|
||
max_bytes: максимальный размер PNG-байт после перекодирования.
|
||
max_side: максимальная сторона в пикселях.
|
||
|
||
Raises:
|
||
ValueError: с человеко-читаемой причиной на русском.
|
||
"""
|
||
from PIL import Image, UnidentifiedImageError
|
||
|
||
if image is None:
|
||
raise ValueError("изображение не передано")
|
||
if not isinstance(image, Image.Image):
|
||
raise ValueError(f"ожидался PIL.Image, получен {type(image).__name__}")
|
||
w, h = image.size
|
||
if max(w, h) > max_side:
|
||
raise ValueError(
|
||
f"изображение слишком большое: {w}x{h}, максимум {max_side}x{max_side}"
|
||
)
|
||
if getattr(image, "format", None) and image.format not in ("PNG", "JPEG", "WEBP"):
|
||
raise ValueError(
|
||
f"неподдерживаемый формат: {image.format}; допустимы PNG, JPEG, WEBP"
|
||
)
|
||
# Прикидка размера: для RGB 3 байта/пиксель + служебные.
|
||
approx = (w * h * 4) + (1024 * 64)
|
||
if approx > max_bytes:
|
||
raise ValueError(
|
||
f"изображение превышает {max_bytes // (1024*1024)} МБ после кодирования"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Удобный высокоуровневый helper, чтобы UI не думал о формате messages.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def generate_svg(
|
||
prompt: str,
|
||
*,
|
||
image_b64: str | None = None,
|
||
mode: str = "icon",
|
||
n: int = 1,
|
||
temperature: float = 0.4,
|
||
model: str | None = None,
|
||
system_prompt: str | None = None,
|
||
) -> list[dict]:
|
||
"""Собирает messages, отправляет в LM Studio и возвращает сырые ответы.
|
||
|
||
Эта функция — тонкая обёртка над `chat()`, вынесенная для удобства UI.
|
||
За реальную сборку messages отвечает `prompts.build_messages`.
|
||
|
||
Args:
|
||
prompt: текст промпта пользователя.
|
||
image_b64: data: URL картинки (или None для text-to-SVG).
|
||
mode: "icon" или "illustration" — для тега system, если не задан явно.
|
||
n: число кандидатов.
|
||
temperature: 0..1.5.
|
||
model: имя модели (опц.).
|
||
system_prompt: если задан — используется как system; иначе caller должен
|
||
сам подставить его в messages (эта функция не знает про prompts/).
|
||
|
||
Returns:
|
||
Список словарей `{'svg_text': str, 'raw': str, 'finish_reason': str}`.
|
||
"""
|
||
if not (1 <= len(prompt.strip()) <= 1000):
|
||
raise ValueError("промпт должен быть от 1 до 1000 символов")
|
||
|
||
messages: list[dict] = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
if image_b64:
|
||
user_content: Any = [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": image_b64}},
|
||
]
|
||
else:
|
||
user_content = prompt
|
||
messages.append({"role": "user", "content": user_content})
|
||
|
||
result = chat(
|
||
messages=messages,
|
||
model=model,
|
||
n=n,
|
||
temperature=temperature,
|
||
)
|
||
return [
|
||
{
|
||
"svg_text": raw, # сырой текст; SVG вытащим extract_svg в валидаторе
|
||
"raw": raw,
|
||
"finish_reason": fr,
|
||
}
|
||
for raw, fr in zip(result.raw_texts, result.finish_reasons)
|
||
]
|