omnisvg-lite/lm_client.py

"""Клиент к LM Studio через OpenAI-compatible chat completions.

Использует синхронный httpx, чтобы не тащить openai SDK и не ловить
несовместимости их версий. Поддерживает текстовые и мультимодальные
(text + image_url) сообщения, а также n>1 для batch-генерации кандидатов.
"""

from __future__ import annotations

import base64
import io
import json
import logging
import os
import time
from dataclasses import dataclass, field
from typing import Any, Iterable, Iterator

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)


@dataclass(frozen=True)
class StreamEvent:
    """Событие потокового ответа LM Studio.

    Attributes:
        type: 'delta' — очередной кусочек текста (delta content от сервера);
            'end' — стрим завершён нормально; 'error' — стрим прерван ошибкой.
        content: текст delta (для type='delta') либо '' для остальных.
        usage: usage-блок, который некоторые серверы шлют в последнем чанке
            (либо None, если не пришёл).
        model: фактическое имя модели из ответа сервера.
        finish_reason: 'stop' / 'length' / 'tool_calls' / '' (если ещё не пришёл).
    """

    type: str  # 'delta' | 'end' | 'error'
    content: str = ""
    usage: dict | None = None
    model: str = ""
    finish_reason: str = ""


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 stream_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,
) -> Iterator[StreamEvent]:
    """Шлёт chat completion в LM Studio со stream=True и отдаёт чанки контента.

    Yields:
        StreamEvent(type='delta', content=...) — очередной кусочек текста.
        StreamEvent(type='end', ...) — финальное событие с метаданными
            (model, finish_reason, usage если сервер прислал).

    Особенности:
        - Reasoning-токены qwen3.5 приходят В ОСНОВНОМ `content` (как обычный
          текст), мы их не отделяем — это работа incremental_svg-парсера в UI.
        - OpenAI не поддерживает `n>1` в стриме. Если пользователь передал
          `n>1`, логируем WARNING и идём с `n=1` в payload.
        - На ошибках (HTTP 4xx/5xx, network, timeout, битый SSE) —
          бросает LMStudioUnavailable. (StreamEvent(type='error') зарезервирован
          на будущее, но в текущей реализации ошибки идут через raise.)

    Raises:
        LMStudioUnavailable при сетевых/HTTP/парсинговых ошибках.
    """
    if n < 1:
        raise ValueError(f"n должно быть >= 1, получено {n}")

    if n > 1:
        log.warning(
            "stream_chat: n=%d запрошено, но в stream-режиме OpenAI не "
            "поддерживает n>1 — идём с n=1",
            n,
        )
        effective_n = 1
    else:
        effective_n = 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": effective_n,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": True,
        # LM Studio / qwen3.5 уважают этот флаг, чтобы не слать thinking
        # отдельным reasoning_content-полем (qwen3.5 кладёт рассуждение в
        # основной content, и мы не пытаемся его отделить).
        "chat_template_kwargs": {"enable_thinking": False},
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {key}",
    }

    log.info(
        "LM Studio (stream) → %s model=%s n=%d temp=%.2f timeout=%.0fs",
        url, mdl, effective_n, temperature, timeout,
    )
    started = time.monotonic()
    # Поля, которые аккумулируются по ходу стрима: финальный чанк часто
    # содержит finish_reason/usage, а model сервер может прислать в первом чанке.
    final_model = ""
    final_usage: dict | None = None
    final_finish_reason = ""
    saw_done = False

    try:
        with httpx.Client(timeout=timeout) as client:
            with client.stream("POST", url, json=payload, headers=headers) as resp:
                # HTTP-ошибки — до чтения тела.
                if resp.status_code >= 400:
                    # Сливаем тело для сообщения об ошибке, но не отдаём
                    # его в стрим.
                    body_preview = ""
                    try:
                        body_preview = resp.read().decode("utf-8", errors="replace")[:200]
                    except Exception:  # noqa: BLE001
                        body_preview = "<no body>"
                    if resp.status_code >= 500:
                        raise LMStudioUnavailable(
                            f"LM Studio error: {resp.status_code} {body_preview}"
                        )
                    raise LMStudioUnavailable(
                        f"LM Studio вернул {resp.status_code}: {body_preview}"
                    )

                # Читаем SSE: каждая строка — это `data: <...>` или пустая
                # строка-разделитель. События разделены пустой строкой.
                for raw_line in resp.iter_lines():
                    if not raw_line:
                        continue
                    # SSE-префикс — `data: ` (с пробелом). Без префикса — мусор.
                    if not raw_line.startswith("data:"):
                        # Может быть комментарий (`: ...`) или event/id —
                        # мы их игнорируем.
                        continue
                    payload_str = raw_line[len("data:"):].strip()
                    if payload_str == "[DONE]":
                        saw_done = True
                        break
                    try:
                        chunk = json.loads(payload_str)
                    except ValueError:
                        # Битый SSE — бросаем, как делает обычный chat().
                        raise LMStudioUnavailable(
                            f"LM Studio stream: не-JSON в SSE-чанке: {payload_str[:200]!r}"
                        )

                    # Достаём метаданные из чанка.
                    if isinstance(chunk, dict):
                        if "model" in chunk and chunk["model"]:
                            final_model = str(chunk["model"])
                        if "usage" in chunk and chunk["usage"]:
                            final_usage = chunk["usage"]

                    choices = chunk.get("choices") or []
                    if not choices:
                        # Heartbeat-чанки без choices — пропускаем.
                        continue
                    first = choices[0]
                    delta = first.get("delta") or {}
                    content = delta.get("content")
                    if content:
                        # content может прийти str или list[dict] (мультимодальный
                        # стрим). Склеиваем в строку.
                        if isinstance(content, list):
                            content = "".join(_coerce_text_part(p) for p in content)
                        yield StreamEvent(type="delta", content=str(content))

                    fr = first.get("finish_reason")
                    if fr:
                        final_finish_reason = str(fr)
    except httpx.TimeoutException as exc:
        raise LMStudioUnavailable(
            f"LM Studio stream timeout at {url}: запрос превысил {timeout:.0f}с"
        ) from exc
    except httpx.HTTPError as exc:
        raise LMStudioUnavailable(
            f"LM Studio stream недоступен по адресу {url}: {exc}"
        ) from exc

    elapsed = time.monotonic() - started
    if not saw_done and not final_finish_reason:
        # Стрим оборвался без [DONE] и без finish_reason — считаем это
        # незавершённым. Не бросаем исключение, чтобы UI мог показать
        # частичный текст; помечаем финал как обрезанный.
        log.warning("LM Studio stream: выход без [DONE] (elapsed=%.2fs)", elapsed)

    yield StreamEvent(
        type="end",
        model=final_model or mdl,
        usage=final_usage,
        finish_reason=final_finish_reason,
    )


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)
    ]