Initial commit: OmniSVG-Lite MVP before live-streaming work
- 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
This commit is contained in:
Mavis
+346
@@ -0,0 +1,346 @@
|
||||
"""Клиент к 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)
|
||||
]
|
||||
Reference in New Issue
Block a user