omnisvg-lite/prompts.py

"""Загрузка промпт-шаблонов и сборка messages[] для LM Studio.

Шаблоны лежат в каталоге `prompts/` рядом с этим модулем:

  - system_icon.txt          — system-инструкция для режима icon
  - system_illustration.txt  — system-инструкция для режима illustration
  - few_shot_examples.txt    — текстовый файл с парами USER:/ASSISTANT:,
                                разделёнными строками `===...===`

Мы не выдумываем промпты — берём ровно то, что лежит в файлах. Это явное
требование задачи.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Literal


# Корень проекта: один уровень вверх от этого файла. Надёжнее, чем os.getcwd().
_PROMPTS_DIR = Path(__file__).resolve().parent / "prompts"

Mode = Literal["icon", "illustration"]


def _read(name: str) -> str:
    """Читает текстовый файл из prompts/; падает с понятной ошибкой, если нет."""
    path = _PROMPTS_DIR / name
    if not path.is_file():
        raise FileNotFoundError(f"промпт-шаблон не найден: {path}")
    return path.read_text(encoding="utf-8").strip()


def load_system_prompt(mode: Mode) -> str:
    """Возвращает system-инструкцию для указанного режима.

    Args:
        mode: "icon" или "illustration".

    Returns:
        Содержимое system_icon.txt / system_illustration.txt.

    Raises:
        ValueError: если mode неизвестен.
    """
    if mode == "icon":
        return _read("system_icon.txt")
    if mode == "illustration":
        return _read("system_illustration.txt")
    raise ValueError(f"неизвестный mode: {mode!r}; ожидается 'icon' или 'illustration'")


# Регулярка для разбора few_shot_examples.txt: делим по строкам из `=`.
_FEW_SHOT_SPLIT = re.compile(r"^=+\s*$", re.MULTILINE)
# Маркеры ролей в каждом блоке.
_TURN_RE = re.compile(r"^(USER|ASSISTANT):\s*\n(.*?)(?=(?:^(?:USER|ASSISTANT):\s*$)|\Z)", re.MULTILINE | re.DOTALL)


def load_few_shot() -> list[dict]:
    """Парсит `prompts/few_shot_examples.txt` и возвращает список сообщений.

    Формат файла: блоки, разделённые строками `===...===`. Внутри блока —
    строки `USER:` и `ASSISTANT:`, после каждой маркерной строки идёт
    содержимое до следующей маркерной строки или до конца блока.

    Returns:
        Список `[{role, content}, ...]` в порядке USER/ASSISTANT пар.
        Блоки, где не нашлось ни одной пары, пропускаются.
    """
    text = _read("few_shot_examples.txt")
    blocks = _FEW_SHOT_SPLIT.split(text)
    messages: list[dict] = []
    for block in blocks:
        # Берём только содержимое блока, убираем шапки типа "EXAMPLE N — ..."
        body = block.strip()
        if not body:
            continue
        # Пропускаем шапки, идущие ДО первого USER/ASSISTANT.
        for turn in _TURN_RE.finditer(body):
            role_token = turn.group(1).lower()
            content = turn.group(2).strip()
            messages.append({"role": role_token, "content": content})
    return messages


def _user_text(prompt: str, *, palette: str | None, has_image: bool) -> str:
    """Собирает финальный текст user-turn с учётом палитры и картинки.

    Args:
        prompt: исходный промпт.
        palette: необязательная палитра.
        has_image: True, если в этот же user-turn пойдёт картинка.

    Returns:
        Готовый текст для content[0].type == "text".
    """
    parts: list[str] = []
    if has_image:
        parts.append(
            "Recreate the visual content of the attached image as SVG. "
            "Do not describe, just generate the markup."
        )
    parts.append(prompt.strip())
    if palette:
        parts.append(f"Palette: {palette.strip()}")
    return "\n".join(parts)


@dataclass(frozen=True)
class BuildResult:
    """Результат сборки messages[] — для удобства дебага/тестов."""

    messages: list[dict]
    mode: Mode
    n: int
    temperature: float


def build_messages(
    prompt: str,
    mode: Mode,
    *,
    image_b64: str | None = None,
    palette: str | None = None,
    n: int = 1,
    temperature: float = 0.4,
) -> list[dict]:
    """Собирает полный список messages для LM Studio.

    Структура:
        1. system
        2. few-shot пары (USER/ASSISTANT) — без картинок
        3. финальный USER: текст (+ опц. image_url)

    Args:
        prompt: пользовательский промпт.
        mode: "icon" или "illustration".
        image_b64: data: URL картинки для image-to-SVG.
        palette: опциональная палитра (добавится как отдельная строка).
        n: число кандидатов (сейчас в messages не подставляется — это параметр
            API-запроса; храним в dataclass для симметрии с дизайном).
        temperature: число в [0, 1.5] (тоже идёт в API, не в messages).

    Returns:
        Список сообщений в формате, понятном `lm_client.chat`.
    """
    if not prompt or not prompt.strip():
        raise ValueError("промпт пустой")
    if mode not in ("icon", "illustration"):
        raise ValueError(f"неизвестный mode: {mode!r}")

    messages: list[dict] = [
        {"role": "system", "content": load_system_prompt(mode)},
    ]
    messages.extend(load_few_shot())

    user_text = _user_text(prompt, palette=palette, has_image=bool(image_b64))
    if image_b64:
        content: list[dict] = [
            {"type": "text", "text": user_text},
            {"type": "image_url", "image_url": {"url": image_b64}},
        ]
    else:
        content = user_text
    messages.append({"role": "user", "content": content})

    # Чтобы лишний раз не плодить dataclass — вернём просто список.
    # Поля n/temperature отдаются API-обёртке; сборка messages от них
    # не зависит (мы лишь гарантируем, что они валидны).
    _ = (n, temperature)
    return messages