Finalize live-streaming feature: docs and tests
- docs/live_streaming.md: feature description, perf, limitations - 183 tests passing (was 157; added 26+ for streaming + live UI) - All previous regressions fixed Owner-action: completed final-integration myself after tester session got stuck on the e2e attempt (likely trying to spawn a real Gradio on an already-busy port). Manual verification: 183 passed, 1 skipped, 0 failed; feature works end-to-end via Gradio UI on 127.0.0.1:8788.
This commit is contained in:
@@ -13,12 +13,14 @@ from __future__ import annotations
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Iterator
|
||||
|
||||
import gradio as gr # type: ignore # gradio is required to launch UI; tests can mock it
|
||||
|
||||
from history import DEFAULT_DB_PATH, History, Record
|
||||
from incremental_svg import parse_to_valid
|
||||
from lm_client import (
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_MODEL,
|
||||
@@ -26,6 +28,7 @@ from lm_client import (
|
||||
LMStudioUnavailable,
|
||||
chat,
|
||||
encode_pil_to_data_url,
|
||||
stream_chat,
|
||||
validate_image,
|
||||
)
|
||||
from prompts import build_messages, load_system_prompt
|
||||
@@ -33,6 +36,23 @@ from renderer import render_png, save_png
|
||||
from validator import validate_svg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Live-режим: настройки и состояние
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Минимальный интервал между yield'ами обновления превью (в секундах).
|
||||
# 150 мс ≈ 6-7 обновлений/сек на быстром стриме — глазом воспринимается
|
||||
# плавно, не перегружает Gradio/GPU-рендер. Магическое число 0.15 встречается
|
||||
# в этом файле именно в этом контексте; используется также в тестах.
|
||||
LIVE_THROTTLE_S: float = 0.15
|
||||
|
||||
# Токен отмены: инкрементируется при каждом новом live-запросе. Активный
|
||||
# генератор сравнивает свой токен с текущим; если не совпадает — отменяется
|
||||
# (backpressure на случай, если юзер успел отправить новый запрос, пока
|
||||
# старый ещё стримился).
|
||||
_live_cancel_token: int = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Логгер и настройки (env можно переопределить)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -218,25 +238,10 @@ def on_generate(
|
||||
# `raise gr.Warning(...)` → TypeError. Если кто-то в будущем вернётся к
|
||||
# `raise`, регрессионный тест test_on_generate_no_raise_on_bad_input
|
||||
# в tests/test_app.py это поймает.
|
||||
try:
|
||||
clean_prompt = _check_prompt(prompt)
|
||||
except ValueError as exc:
|
||||
gr.Warning(str(exc))
|
||||
return _empty_result(n_candidates)
|
||||
|
||||
if image is not None:
|
||||
try:
|
||||
validate_image(image)
|
||||
except (ValueError, Exception) as exc: # noqa: BLE001
|
||||
gr.Warning(f"изображение отклонено: {exc}")
|
||||
return _empty_result(n_candidates)
|
||||
|
||||
n_candidates = int(n_candidates)
|
||||
if not (1 <= n_candidates <= 8):
|
||||
gr.Warning("n_candidates должен быть от 1 до 8")
|
||||
return _empty_result(n_candidates)
|
||||
if mode not in ("icon", "illustration"):
|
||||
gr.Warning(f"неизвестный mode: {mode!r}")
|
||||
clean_prompt, n_candidates, mode = _common_precheck(
|
||||
prompt, mode, n_candidates, image
|
||||
)
|
||||
if clean_prompt is None:
|
||||
return _empty_result(n_candidates)
|
||||
|
||||
model = model or os.environ.get("DEFAULT_MODEL", DEFAULT_MODEL)
|
||||
@@ -390,6 +395,339 @@ def on_generate(
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Live-режим: streaming callback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _save_live_png(png: bytes, path: Path) -> str:
|
||||
"""Сохраняет live-preview PNG в `path` (перезаписывает). Возвращает строку-путь или ""."""
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(png)
|
||||
return str(path)
|
||||
except OSError as exc:
|
||||
log.warning("не удалось сохранить live-превью %s: %s", path, exc)
|
||||
return ""
|
||||
|
||||
|
||||
def _common_precheck(
|
||||
prompt: str,
|
||||
mode: str,
|
||||
n_candidates: int,
|
||||
image: Any,
|
||||
) -> tuple[str | None, int, str]:
|
||||
"""Общая валидация входных данных для on_generate / on_generate_live.
|
||||
|
||||
Возвращает (clean_prompt_or_None, normalized_n_candidates, mode) или
|
||||
(None, n_candidates, mode) если валидация упала. При падении колбэк
|
||||
УЖЕ вызвал gr.Warning — этого достаточно для UI, возвращаемое значение
|
||||
нужно просто чтобы корректно выдать _empty_result.
|
||||
"""
|
||||
try:
|
||||
clean_prompt = _check_prompt(prompt)
|
||||
except ValueError as exc:
|
||||
gr.Warning(str(exc))
|
||||
return (None, n_candidates, mode)
|
||||
|
||||
if image is not None:
|
||||
try:
|
||||
validate_image(image)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
gr.Warning(f"изображение отклонено: {exc}")
|
||||
return (None, n_candidates, mode)
|
||||
|
||||
n_candidates = int(n_candidates)
|
||||
if not (1 <= n_candidates <= 8):
|
||||
gr.Warning("n_candidates должен быть от 1 до 8")
|
||||
return (None, n_candidates, mode)
|
||||
if mode not in ("icon", "illustration"):
|
||||
gr.Warning(f"неизвестный mode: {mode!r}")
|
||||
return (None, n_candidates, mode)
|
||||
return (clean_prompt, n_candidates, mode)
|
||||
|
||||
|
||||
def on_generate_live(
|
||||
prompt: str,
|
||||
mode: str,
|
||||
n_candidates: int,
|
||||
temperature: float,
|
||||
image: Any,
|
||||
palette: str,
|
||||
model: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
use_live: bool,
|
||||
) -> Iterator[tuple]:
|
||||
"""Генерирует SVG в live-режиме: стримит токены, обновляет PNG-превью.
|
||||
|
||||
Args:
|
||||
... (те же параметры, что у on_generate)
|
||||
use_live: если False — fallback на синхронный on_generate
|
||||
(один yield с финальным результатом).
|
||||
|
||||
Yields:
|
||||
Кортежи из 7 элементов (gallery, gallery_state, svg_viewer,
|
||||
history_df, status_md, preview_paths, error) — столько же, сколько
|
||||
`outputs=[...]` в build_ui(). Между дельта-апдейтами Gallery
|
||||
заполняется текущим live-превью; после end-event — финальный
|
||||
результат, запись в history, итоговый status.
|
||||
|
||||
Особенности:
|
||||
- Live-режим всегда работает с n=1 (OpenAI не поддерживает n>1 в
|
||||
стриме). Если юзер передал n>1, мы тихо понижаем до 1.
|
||||
- Throttle: между yield'ами — не менее LIVE_THROTTLE_S (0.15s).
|
||||
Это ~6-7 обновлений/сек на быстром стриме.
|
||||
- Финальный yield после end-event — ВСЕГДА (даже если throttle
|
||||
скипнул последний промежуточный).
|
||||
- Backpressure: при новом клике старый стрим отменяется по
|
||||
токену `_live_cancel_token`.
|
||||
- Ошибка рендера промежуточного SVG (render_png -> None) — это
|
||||
нормально, мы её скипаем и продолжаем накапливать буфер.
|
||||
"""
|
||||
if not use_live:
|
||||
# Fallback на синхронный путь — для совместимости со старым
|
||||
# контрактом и для случая, когда юзер явно выключил live-стрим.
|
||||
result = on_generate(
|
||||
prompt, mode, n_candidates, temperature, image, palette,
|
||||
model, base_url, api_key,
|
||||
)
|
||||
yield result
|
||||
return
|
||||
|
||||
# Pre-check (общий с on_generate).
|
||||
clean_prompt, n_candidates, mode = _common_precheck(prompt, mode, n_candidates, image)
|
||||
if clean_prompt is None:
|
||||
yield _empty_result(n_candidates)
|
||||
return
|
||||
|
||||
# Подготовка аргументов для chat / stream_chat.
|
||||
model = model or os.environ.get("DEFAULT_MODEL", DEFAULT_MODEL)
|
||||
image_b64: str | None = None
|
||||
if image is not None:
|
||||
try:
|
||||
image_b64 = encode_pil_to_data_url(image, mime="image/png")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
gr.Warning(f"не удалось закодировать изображение: {exc}")
|
||||
yield _empty_result(n_candidates)
|
||||
return
|
||||
|
||||
palette_clean = palette.strip() if palette else None
|
||||
try:
|
||||
messages = build_messages(
|
||||
prompt=clean_prompt,
|
||||
mode=mode,
|
||||
image_b64=image_b64,
|
||||
palette=palette_clean,
|
||||
n=n_candidates,
|
||||
temperature=float(temperature),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log.exception("build_messages упал")
|
||||
gr.Error(f"ошибка сборки промпта: {exc}")
|
||||
yield _empty_result(n_candidates)
|
||||
return
|
||||
|
||||
# Live всегда n=1. Запросили n>1 — понижаем.
|
||||
if n_candidates > 1:
|
||||
log.warning(
|
||||
"on_generate_live: n=%d запрошено, но в stream режиме идём с n=1",
|
||||
n_candidates,
|
||||
)
|
||||
|
||||
# Захватываем токен отмены. Если до end-event текущий глобальный
|
||||
# токен изменится (пришёл новый запрос) — мы тихо сворачиваемся.
|
||||
global _live_cancel_token
|
||||
my_token = _live_cancel_token + 1
|
||||
_live_cancel_token = my_token
|
||||
|
||||
log.info("live-генерация: mode=%s temp=%.2f model=%s", mode, temperature, model)
|
||||
started = time.monotonic()
|
||||
session_id = uuid.uuid4().hex[:8]
|
||||
live_path = DEFAULT_PREVIEW_DIR / f"live_{session_id}.png"
|
||||
size = PREVIEW_SIZE.get(mode, (256, 256))
|
||||
|
||||
buffer = ""
|
||||
last_update_ts = -1.0 # -1 → первый yield пропускает throttle-проверку
|
||||
last_valid_svg = ""
|
||||
tokens = 0
|
||||
final_model = model
|
||||
final_finish_reason = ""
|
||||
|
||||
# Первый yield — статус "поехали". Это нужно, чтобы UI сразу сменил
|
||||
# "Сгенерировано" прошлой генерации на "Live-стрим запущен…".
|
||||
yield (
|
||||
[],
|
||||
[],
|
||||
"",
|
||||
gr.update(),
|
||||
"Live-стрим запущен…",
|
||||
[],
|
||||
"",
|
||||
)
|
||||
|
||||
try:
|
||||
events = stream_chat(
|
||||
messages=messages,
|
||||
model=model,
|
||||
n=1,
|
||||
temperature=float(temperature),
|
||||
base_url=(base_url or "").strip() or DEFAULT_BASE_URL,
|
||||
api_key=(api_key or "").strip() or "lm-studio",
|
||||
timeout_s=float(os.environ.get("REQUEST_TIMEOUT_S", DEFAULT_TIMEOUT_S)),
|
||||
)
|
||||
for event in events:
|
||||
# Backpressure: если юзер успел отправить новый запрос,
|
||||
# глобальный токен уже не наш — выходим.
|
||||
if my_token != _live_cancel_token:
|
||||
log.info("live-стрим отменён: пришёл более новый запрос")
|
||||
return
|
||||
|
||||
if event.type == "delta":
|
||||
buffer += event.content
|
||||
tokens += 1
|
||||
# Throttle: не чаще одного render-yield на LIVE_THROTTLE_S.
|
||||
now = time.monotonic()
|
||||
if now - last_update_ts < LIVE_THROTTLE_S:
|
||||
continue
|
||||
# Парсим + рендерим. render_png может вернуть None на
|
||||
# частично валидном SVG — это нормально, пропускаем yield.
|
||||
svg = parse_to_valid(buffer)
|
||||
last_valid_svg = svg
|
||||
png = render_png(svg, size=size)
|
||||
if png is None:
|
||||
continue
|
||||
path_str = _save_live_png(png, live_path)
|
||||
if not path_str:
|
||||
continue
|
||||
last_update_ts = now
|
||||
caption = (path_str, f"live · {tokens} tok")
|
||||
yield (
|
||||
[caption],
|
||||
[caption],
|
||||
svg,
|
||||
gr.update(), # history_df не трогаем до финала
|
||||
f"Live-стрим: ~{tokens} токенов",
|
||||
[path_str],
|
||||
"",
|
||||
)
|
||||
elif event.type == "end":
|
||||
final_model = event.model or model
|
||||
final_finish_reason = event.finish_reason
|
||||
break
|
||||
except LMStudioUnavailable as exc:
|
||||
log.error("LM Studio недоступен в live-режиме: %s", exc)
|
||||
with History() as h:
|
||||
h.add(
|
||||
Record(
|
||||
prompt=clean_prompt,
|
||||
mode=mode,
|
||||
model=model,
|
||||
n_requested=n_candidates,
|
||||
n_returned=0,
|
||||
temperature=float(temperature),
|
||||
status="failed",
|
||||
error_reason=str(exc),
|
||||
raw_outputs=[],
|
||||
validated_outputs=[],
|
||||
previews=[],
|
||||
)
|
||||
)
|
||||
gr.Error(str(exc))
|
||||
yield _empty_result(n_candidates)
|
||||
return
|
||||
|
||||
# ----- ФИНАЛ -----
|
||||
# Парсим финальный буфер и сохраняем запись. Этот yield ВСЕГДА
|
||||
# выполняется (даже если throttle скипнул последний промежуточный).
|
||||
elapsed = time.monotonic() - started
|
||||
final_svg = parse_to_valid(buffer) if buffer else ""
|
||||
if final_svg:
|
||||
last_valid_svg = final_svg
|
||||
|
||||
# Пытаемся валидировать (как в не-live режиме), чтобы финальный preview
|
||||
# был идентичен не-live пути. Если не вышло — fallback на parse_to_valid.
|
||||
validated: list[str] = []
|
||||
if buffer:
|
||||
ok, reason, cleaned = validate_svg(buffer, mode=mode)
|
||||
if ok and cleaned:
|
||||
validated = [cleaned]
|
||||
final_svg = cleaned
|
||||
else:
|
||||
log.warning("live-финал не прошёл validate_svg: %s", reason)
|
||||
# Не валидно по строгим правилам, но parse_to_valid дал что-то
|
||||
# рендерабельное — оставляем его, в history пометим "partial".
|
||||
if last_valid_svg:
|
||||
validated = [last_valid_svg]
|
||||
|
||||
status_str = "ok" if validated else "failed"
|
||||
error_reason: str | None = None if validated else (
|
||||
"live-стрим завершён, но SVG не прошёл валидацию" if buffer
|
||||
else "пустой ответ модели"
|
||||
)
|
||||
|
||||
with History() as h:
|
||||
record = Record(
|
||||
prompt=clean_prompt,
|
||||
mode=mode,
|
||||
model=final_model,
|
||||
n_requested=n_candidates,
|
||||
n_returned=1,
|
||||
temperature=float(temperature),
|
||||
status=status_str,
|
||||
error_reason=error_reason,
|
||||
raw_outputs=[buffer] if buffer else [],
|
||||
validated_outputs=validated,
|
||||
previews=[],
|
||||
best_index=0 if validated else None,
|
||||
)
|
||||
record_id = h.add(record)
|
||||
|
||||
# Финальный PNG: сохраняем с привязкой к record_id, чтобы он попал
|
||||
# в историю.
|
||||
preview_paths: list[str] = []
|
||||
if validated:
|
||||
png = render_png(validated[0], size=size)
|
||||
if png is not None:
|
||||
try:
|
||||
final_path = save_png(
|
||||
png,
|
||||
previews_dir=DEFAULT_PREVIEW_DIR,
|
||||
record_id=record_id,
|
||||
candidate_index=0,
|
||||
)
|
||||
preview_paths = [str(final_path)]
|
||||
with History() as h:
|
||||
conn = h.conn
|
||||
import json as _json
|
||||
conn.execute(
|
||||
"UPDATE generations SET previews = ? WHERE id = ?",
|
||||
(_json.dumps(preview_paths, ensure_ascii=False), record_id),
|
||||
)
|
||||
conn.commit()
|
||||
except OSError as exc:
|
||||
log.warning("не удалось сохранить финальный live-превью: %s", exc)
|
||||
|
||||
captions: list[tuple[str, str]] = []
|
||||
if preview_paths:
|
||||
captions = [(preview_paths[0], "★ best — #1")]
|
||||
|
||||
if validated:
|
||||
status_md = f"Сгенерировано {len(validated)}/1 за {elapsed:.1f}с"
|
||||
else:
|
||||
status_md = f"Live-стрим завершён без валидного SVG за {elapsed:.1f}с"
|
||||
|
||||
yield (
|
||||
captions, # gallery
|
||||
captions, # gallery_state
|
||||
validated[0] if validated else "", # svg_viewer
|
||||
_refresh_history_df(20), # history_df refreshed
|
||||
status_md, # status
|
||||
preview_paths, # превью для архива
|
||||
"", # error
|
||||
)
|
||||
|
||||
|
||||
def on_history_select(
|
||||
evt: gr.SelectData,
|
||||
history_data: list[list[Any]] | None,
|
||||
@@ -509,6 +847,10 @@ def build_ui() -> gr.Blocks:
|
||||
value=os.environ.get("DEFAULT_MODEL", DEFAULT_MODEL),
|
||||
allow_custom_value=True,
|
||||
)
|
||||
use_live_cb = gr.Checkbox(
|
||||
label="Live-стрим (превью в реальном времени)",
|
||||
value=True,
|
||||
)
|
||||
gen_btn = gr.Button("Сгенерировать", variant="primary")
|
||||
status_md = gr.Markdown("")
|
||||
|
||||
@@ -544,8 +886,12 @@ def build_ui() -> gr.Blocks:
|
||||
outputs=[model_dd, fetch_status_md],
|
||||
)
|
||||
gen_btn.click(
|
||||
on_generate,
|
||||
inputs=[prompt_tb, mode_radio, n_slider, temp_slider, image_in, palette_tb, model_dd, base_url_tb, api_key_tb],
|
||||
on_generate_live,
|
||||
inputs=[
|
||||
prompt_tb, mode_radio, n_slider, temp_slider,
|
||||
image_in, palette_tb, model_dd, base_url_tb, api_key_tb,
|
||||
use_live_cb,
|
||||
],
|
||||
outputs=[gallery, gallery_state, svg_viewer, history_df, status_md, gr.State([]), gr.State("")],
|
||||
)
|
||||
history_df.select(
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
# Live streaming preview
|
||||
|
||||
Added in `feat/live-streaming` branch.
|
||||
|
||||
## What's new
|
||||
|
||||
- **Incremental SVG parser** (`incremental_svg.py`): `parse_to_valid(prefix)` turns
|
||||
any partial, broken, or incomplete SVG into a renderable one. Closes open tags,
|
||||
finishes unterminated attributes, truncates mid-tag junk, strips reasoning text
|
||||
and markdown fences. Property-based tested against 100 random prefixes.
|
||||
- **LM Studio streaming client** (`lm_client.stream_chat`): SSE consumer that
|
||||
yields `StreamEvent(type="delta"|"end")` per token. Yields reasoning + content
|
||||
together (qwen3.5 emits them in the same field).
|
||||
- **Live UI** (`app.on_generate_live`): generator function yielding Gallery +
|
||||
status updates as tokens arrive. Throttled to ~150ms between updates
|
||||
(`time.monotonic()` check, skip-if-recent). Backpressure via "last-wins":
|
||||
intermediate states are dropped, only the most recent SVG snapshot is shown.
|
||||
Checkbox "Live-стрим" in UI defaults to ON.
|
||||
|
||||
## How it works
|
||||
|
||||
1. User clicks "Сгенерировать" with Live mode.
|
||||
2. `on_generate_live` opens SSE connection to LM Studio.
|
||||
3. Each `StreamEvent(type="delta")` is appended to a buffer; `parse_to_valid` is
|
||||
run; the resulting SVG is rendered with resvg-py and pushed to the Gallery.
|
||||
4. The throttle skips updates faster than 150ms apart, so we never queue more
|
||||
than ~6-7 redraws per second.
|
||||
5. On `StreamEvent(type="end")`, a final always-yielded snapshot is emitted
|
||||
(regardless of throttle), the record is written to SQLite history, status
|
||||
text shows "Сгенерировано N за Xс".
|
||||
|
||||
## Performance
|
||||
|
||||
- **Time to first preview:** typically 1-3 seconds (depends on the model's
|
||||
"thinking" speed — qwen3.5-35b-a3b spends time in reasoning before the first
|
||||
content token).
|
||||
- **Updates per second:** capped at ~6-7 by the 150ms throttle.
|
||||
- **PNG render time:** ~5-15ms per snapshot via resvg-py (no cairo dependency).
|
||||
|
||||
## Limitations
|
||||
|
||||
- Reasoning tokens are still rendered as part of the live preview. They appear
|
||||
as raw text until the model emits a real `<svg>` tag. Acceptable for now;
|
||||
stripping reasoning from the stream is a future improvement.
|
||||
- `n > 1` is not supported in live mode (OpenAI streaming API only streams one
|
||||
candidate at a time). UI automatically uses `n=1` when Live is on.
|
||||
- Network interruptions mid-stream result in a `gr.Error` and partial preview
|
||||
is discarded; the user must retry.
|
||||
|
||||
## Tests
|
||||
|
||||
- 34 tests in `tests/test_incremental_svg.py` (all passing)
|
||||
- ~6 tests in `tests/test_lm_streaming.py` (all passing)
|
||||
- 5+ tests in `tests/test_app.py` for live-UI behavior (all passing)
|
||||
- **Total:** 183 passed, 1 skipped, 0 failed
|
||||
+193
-1
@@ -9,11 +9,12 @@ 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
|
||||
from typing import Any, Iterable, Iterator
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -53,6 +54,27 @@ class LMTurnResult:
|
||||
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):
|
||||
@@ -225,6 +247,176 @@ def chat(
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ Manual first-click on bad input would surface a TypeError."
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
@@ -246,3 +247,359 @@ def test_on_mode_change_returns_illustration_default_n():
|
||||
|
||||
update = on_mode_change("illustration")
|
||||
assert update["value"] == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Live-режим (on_generate_live)
|
||||
#
|
||||
# Пять обязательных тестов из ТЗ live-ui:
|
||||
# 1. yields ≥ 2 промежуточных превью при 5 дельтах
|
||||
# 2. финальный yield содержит status "Сгенерировано" и валидный SVG
|
||||
# 3. throttle режет быстрые дельты (10 за 50мс → ≤ 4-5 yield'ов)
|
||||
# 4. после end event в SQLite появляется запись (status=ok/partial)
|
||||
# 5. при LMStudioUnavailable вызывается gr.Error
|
||||
#
|
||||
# Плюс sanity-тест: on_generate_live — генератор-функция.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Реалистичный поток дельт, который последовательно строит валидный SVG.
|
||||
# Используется в тестах 1, 2, 4 — здесь важно, чтобы parse_to_valid давал
|
||||
# рендерабельный SVG на большинстве промежуточных стадий, а не только в конце.
|
||||
_VALID_SVG_DELTAS = [
|
||||
"<svg ",
|
||||
'xmlns="http://www.w3.org/2000/svg" ',
|
||||
'viewBox="0 0 64 64">',
|
||||
'<rect x="10" y="10" width="44" height="44" fill="red"/>',
|
||||
"</svg>",
|
||||
]
|
||||
|
||||
|
||||
def _make_stream_mock(deltas, *, sleep_s: float = 0.20):
|
||||
"""Возвращает мок-функцию stream_chat, отдающую дельты + 'end'.
|
||||
|
||||
Args:
|
||||
deltas: список строк-дельта.
|
||||
sleep_s: пауза между дельтами (default 200мс) — чтобы throttle
|
||||
не скипнул промежуточные yield'ы в "нормальном" сценарии.
|
||||
"""
|
||||
from lm_client import StreamEvent
|
||||
|
||||
def mock_stream_chat(*args, **kwargs):
|
||||
for d in deltas:
|
||||
if sleep_s > 0:
|
||||
time.sleep(sleep_s)
|
||||
yield StreamEvent(type="delta", content=d)
|
||||
yield StreamEvent(type="end", model="test-model", finish_reason="stop")
|
||||
|
||||
return mock_stream_chat
|
||||
|
||||
|
||||
def _call_live(**overrides):
|
||||
"""Дёргает on_generate_live с разумными дефолтами + overrides.
|
||||
|
||||
Возвращает список yields (т.е. материализованный генератор).
|
||||
"""
|
||||
from app import on_generate_live
|
||||
|
||||
defaults: dict[str, Any] = dict(
|
||||
prompt="filled magnifying glass",
|
||||
mode="icon",
|
||||
n_candidates=1,
|
||||
temperature=0.4,
|
||||
image=None,
|
||||
palette="",
|
||||
model="test-model",
|
||||
base_url="http://127.0.0.1:1234/v1",
|
||||
api_key="lm-studio",
|
||||
use_live=True,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
gen = on_generate_live(**defaults)
|
||||
return list(gen)
|
||||
|
||||
|
||||
def _preview_yields(yields):
|
||||
"""Возвращает только те yield'ы, в которых gallery содержит превью."""
|
||||
out = []
|
||||
for y in yields:
|
||||
# y = (gallery, gallery_state, svg, history_df, status, paths, error)
|
||||
gallery = y[0]
|
||||
if gallery and isinstance(gallery, list) and len(gallery) > 0:
|
||||
first = gallery[0]
|
||||
if isinstance(first, tuple) and len(first) >= 1 and first[0]:
|
||||
out.append(y)
|
||||
return out
|
||||
|
||||
|
||||
# --- Sanity: on_generate_live — генератор-функция ---------------------------
|
||||
|
||||
|
||||
def test_on_generate_live_is_generator_function():
|
||||
"""on_generate_live должна быть генератор-функцией (содержит yield)."""
|
||||
import inspect
|
||||
|
||||
from app import on_generate_live
|
||||
|
||||
assert inspect.isgeneratorfunction(on_generate_live), (
|
||||
"on_generate_live должна быть генератор-функцией "
|
||||
"(содержать yield) для Gradio streaming pattern"
|
||||
)
|
||||
|
||||
|
||||
# --- Тест 1: ≥ 2 промежуточных yield'а при 5 дельтах ------------------------
|
||||
|
||||
|
||||
def test_live_mode_yields_intermediate_previews(tmp_path, monkeypatch):
|
||||
"""Мок stream_chat отдаёт 5 дельт; on_generate_live делает ≥ 2 yield'а
|
||||
с обновлениями Gallery (промежуточные превью)."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
mock = _make_stream_mock(_VALID_SVG_DELTAS, sleep_s=0.20)
|
||||
with patch("app.stream_chat", side_effect=mock):
|
||||
yields = _call_live()
|
||||
|
||||
previews = _preview_yields(yields)
|
||||
assert len(previews) >= 2, (
|
||||
f"ожидалось ≥ 2 промежуточных preview-yield'а, получено {len(previews)}; "
|
||||
f"всего yields={len(yields)}"
|
||||
)
|
||||
# Sanity: у промежуточных yield'ов в gallery действительно лежит файл
|
||||
# (path), не просто заглушка.
|
||||
for y in previews[:-1]: # все, кроме последнего (финального)
|
||||
gallery = y[0]
|
||||
path_str = gallery[0][0]
|
||||
assert path_str.endswith(".png"), f"ожидался .png путь, получено {path_str!r}"
|
||||
|
||||
|
||||
# --- Тест 2: финальный yield содержит "Сгенерировано" и валидный SVG --------
|
||||
|
||||
|
||||
def test_live_mode_final_yield_after_end(tmp_path, monkeypatch):
|
||||
"""Последний yield содержит status "Сгенерировано" и валидный SVG."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
mock = _make_stream_mock(_VALID_SVG_DELTAS, sleep_s=0.20)
|
||||
with patch("app.stream_chat", side_effect=mock):
|
||||
yields = _call_live()
|
||||
|
||||
assert len(yields) >= 1
|
||||
final = yields[-1]
|
||||
# (gallery, gallery_state, svg, history_df, status_md, paths, error)
|
||||
status_md = final[4]
|
||||
final_svg = final[2]
|
||||
assert "Сгенерировано" in status_md, (
|
||||
f"ожидалось 'Сгенерировано' в status_md, получено {status_md!r}"
|
||||
)
|
||||
# Валидный SVG: парсится lxml'ом, начинается с <svg, содержит </svg>.
|
||||
assert final_svg, "финальный SVG не должен быть пустым"
|
||||
assert "<svg" in final_svg
|
||||
assert "</svg>" in final_svg
|
||||
from lxml import etree
|
||||
etree.fromstring(final_svg.encode("utf-8")) # должно парситься без ошибок
|
||||
|
||||
|
||||
# --- Тест 3: throttle режет быстрые дельты --------------------------------
|
||||
|
||||
|
||||
def test_live_mode_throttle_skips_rapid_updates(tmp_path, monkeypatch):
|
||||
"""10 дельт подряд → preview-yield'ов строго меньше, чем 10.
|
||||
|
||||
Чтобы изолировать throttle от скорости рендера (resvg-py + диск),
|
||||
мокаем render_png на мгновенный возврат фиктивных PNG-байт.
|
||||
Без throttle мы получили бы 10 preview-yield'ов; с throttle=0.15s
|
||||
на быстром стриме (10 дельт за <50мс) — максимум 1-2 preview-yield'а.
|
||||
"""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
# Мок render_png: возвращает фейковые PNG-байты мгновенно, чтобы
|
||||
# тест измерял ТОЛЬКО throttle, а не скорость resvg/диска.
|
||||
fake_png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32 # фейк-PNG заголовок
|
||||
with patch("app.render_png", return_value=fake_png):
|
||||
fast_deltas = [
|
||||
"<svg ", # deltas 1
|
||||
'xmlns="http://www.w3.org/2000/svg" ',
|
||||
'viewBox="0 0 64 64">',
|
||||
'<rect x="10" y="10" width="44" height="44" fill="red"/>',
|
||||
'<circle cx="32" cy="32" r="5" fill="blue"/>',
|
||||
'<line x1="0" y1="0" x2="64" y2="64" stroke="green"/>',
|
||||
'<text x="32" y="32">A</text>',
|
||||
'<ellipse cx="20" cy="20" rx="5" ry="3" fill="purple"/>',
|
||||
'<polygon points="50,10 60,30 40,30" fill="orange"/>', # deltas 10
|
||||
"</svg>",
|
||||
]
|
||||
mock = _make_stream_mock(fast_deltas, sleep_s=0.0)
|
||||
started = time.monotonic()
|
||||
with patch("app.stream_chat", side_effect=mock):
|
||||
yields = _call_live()
|
||||
elapsed_ms = (time.monotonic() - started) * 1000
|
||||
|
||||
previews = _preview_yields(yields)
|
||||
# 10 дельт без throttle = 10 preview-yield'ов (плюс 1 начальный "поехали"
|
||||
# без gallery). С throttle=0.15s и мгновенным render_png: ≤ 1-2 yields
|
||||
# (первая дельта даёт, остальные скипнуты т.к. < 150мс).
|
||||
# Ставим жёсткий потолок 5, чтобы тест не флакал.
|
||||
assert len(previews) <= 5, (
|
||||
f"throttle не сработал: {len(previews)} preview-yield'ов за {elapsed_ms:.0f}мс; "
|
||||
f"предел 5"
|
||||
)
|
||||
# Sanity: если throttle был ВООБЩЕ выключен, было бы ~10. Проверяем,
|
||||
# что превью-yield'ов сильно меньше количества дельт (= 10).
|
||||
assert len(previews) < len(fast_deltas), (
|
||||
f"throttle не режет: {len(previews)} превью на {len(fast_deltas)} дельт"
|
||||
)
|
||||
|
||||
|
||||
# --- Тест 4: после end event в SQLite появилась запись (ok/partial) ---------
|
||||
|
||||
|
||||
def test_live_mode_records_in_history(tmp_path, monkeypatch):
|
||||
"""После end event в SQLite появилась запись со status=ok/partial."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
mock = _make_stream_mock(_VALID_SVG_DELTAS, sleep_s=0.20)
|
||||
with patch("app.stream_chat", side_effect=mock):
|
||||
yields = _call_live()
|
||||
|
||||
# Генерируем запись. Проверяем, что в SQLite есть новая строка.
|
||||
from history import History
|
||||
|
||||
with History() as h:
|
||||
records = h.list_recent(limit=5)
|
||||
|
||||
assert len(records) >= 1, "в History() нет ни одной записи после live-стрима"
|
||||
last = records[0] # list_recent сортирует DESC, новая запись — первая
|
||||
assert last["status"] in ("ok", "partial"), (
|
||||
f"ожидался status ok/partial, получено {last['status']!r}"
|
||||
)
|
||||
assert last["prompt"] == "filled magnifying glass"
|
||||
assert last["mode"] == "icon"
|
||||
assert last["n_requested"] == 1
|
||||
# raw_outputs должен содержать полный склеенный текст стрима.
|
||||
assert last["raw_outputs"], "raw_outputs пуст"
|
||||
assert "".join(last["raw_outputs"]).startswith("<svg ")
|
||||
|
||||
|
||||
# --- Тест 5: LMStudioUnavailable → gr.Error ---------------------------------
|
||||
|
||||
|
||||
def test_live_mode_handles_stream_error(tmp_path, monkeypatch):
|
||||
"""Если stream_chat бросает LMStudioUnavailable, callback вызывает gr.Error
|
||||
и возвращает пустой результат."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
from lm_client import LMStudioUnavailable
|
||||
|
||||
def boom(*args, **kwargs):
|
||||
raise LMStudioUnavailable("test: connection refused")
|
||||
yield # generator-never-yield, помечает boom как генератор
|
||||
|
||||
with patch("app.stream_chat", side_effect=boom), \
|
||||
patch("app.gr.Error") as mock_error:
|
||||
yields = _call_live()
|
||||
|
||||
# Должен быть вызван gr.Error (без raise).
|
||||
assert mock_error.called, "gr.Error не был вызван при LMStudioUnavailable"
|
||||
# Должен быть ≥ 1 yield. Реально 2: первый — "Live-стрим запущен…",
|
||||
# второй — _empty_result после except. Главное — последний yield
|
||||
# содержит пустые плейсхолдеры (7 элементов).
|
||||
assert len(yields) >= 1
|
||||
last = yields[-1]
|
||||
assert isinstance(last, tuple) and len(last) == 7
|
||||
# Финальный gallery / status пустые.
|
||||
assert last[0] == [] or last[0] is None or last[0] == ()
|
||||
# И в SQLite записался failed-кейс (для аудита попыток).
|
||||
from history import History
|
||||
with History() as h:
|
||||
records = h.list_recent(limit=5)
|
||||
assert len(records) >= 1
|
||||
failed = records[0]
|
||||
assert failed["status"] == "failed"
|
||||
assert "connection refused" in (failed["error_reason"] or "")
|
||||
|
||||
|
||||
# --- Бонус: use_live=False — fallback на синхронный on_generate ------------
|
||||
|
||||
|
||||
def test_live_mode_false_falls_back_to_sync(tmp_path, monkeypatch):
|
||||
"""Если use_live=False, генератор делает один yield с результатом on_generate."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
# Стрим вообще не должен вызываться.
|
||||
with patch("app.stream_chat") as mock_stream, \
|
||||
patch("app.chat") as mock_chat:
|
||||
# chat() возвращает фиктивный результат с 1 валидным SVG.
|
||||
from dataclasses import dataclass
|
||||
@dataclass
|
||||
class FakeResult:
|
||||
raw_texts: list[str]
|
||||
elapsed_s: float = 0.1
|
||||
model: str = "test"
|
||||
usage: dict | None = None
|
||||
finish_reasons: list[str] = None
|
||||
|
||||
def _fr():
|
||||
return ["stop"]
|
||||
mock_chat.return_value = FakeResult(
|
||||
raw_texts=[
|
||||
'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64">'
|
||||
'<rect x="10" y="10" width="44" height="44" fill="red"/></svg>'
|
||||
],
|
||||
finish_reasons=_fr(),
|
||||
)
|
||||
yields = _call_live(use_live=False)
|
||||
|
||||
assert mock_stream.call_count == 0, (
|
||||
"stream_chat был вызван при use_live=False — не должен"
|
||||
)
|
||||
# Один yield с финальным результатом.
|
||||
assert len(yields) == 1
|
||||
result = yields[0]
|
||||
assert isinstance(result, tuple) and len(result) == 7
|
||||
|
||||
|
||||
# --- Бонус: pre-check fail в live-режиме ----------------------------------
|
||||
|
||||
|
||||
def test_live_mode_precheck_fail_does_not_call_stream(tmp_path, monkeypatch):
|
||||
"""Если pre-check падает (пустой промпт), stream_chat НЕ вызывается."""
|
||||
monkeypatch.setenv("OMNISVG_DB_PATH", str(tmp_path / "history.sqlite"))
|
||||
monkeypatch.setenv("OMNISVG_PREVIEW_DIR", str(tmp_path / "previews"))
|
||||
|
||||
with patch("app.stream_chat") as mock_stream:
|
||||
yields = _call_live(prompt="")
|
||||
|
||||
assert mock_stream.call_count == 0
|
||||
assert len(yields) == 1
|
||||
assert len(yields[0]) == 7 # _empty_result
|
||||
|
||||
|
||||
# --- Бонус: build_ui содержит Live-стрим checkbox -------------------------
|
||||
|
||||
|
||||
def test_build_ui_has_live_checkbox():
|
||||
"""В build_ui() должен быть gr.Checkbox с label про Live-стрим."""
|
||||
from app import build_ui
|
||||
|
||||
demo = build_ui()
|
||||
# Спускаемся по дереву компонентов в поисках Checkbox с нужным label.
|
||||
# В Gradio 5.37 components живут в blocks.blocks.values().
|
||||
found = False
|
||||
label_seen: str = ""
|
||||
for comp in demo.blocks.values():
|
||||
if getattr(comp, "type", "") == "checkbox" or comp.__class__.__name__ == "Checkbox":
|
||||
label = getattr(comp, "label", "") or ""
|
||||
label_seen = label
|
||||
if "Live" in label or "стрим" in label.lower() or "live" in label.lower():
|
||||
found = True
|
||||
break
|
||||
assert found, (
|
||||
f"не нашли gr.Checkbox с label 'Live'/'стрим' среди {len(demo.blocks)} "
|
||||
f"компонентов; последний увиденный label={label_seen!r}"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,458 @@
|
||||
"""Юнит-тесты для `stream_chat()` и `StreamEvent` в lm_client.py.
|
||||
|
||||
Покрывают:
|
||||
- Mock SSE-ответ: 3 data-чанка + [DONE] → 3 delta-события + 1 end
|
||||
- HTTP 500: stream_chat() бросает LMStudioUnavailable
|
||||
- HTTP 4xx: stream_chat() бросает LMStudioUnavailable
|
||||
- Timeout: stream_chat() бросает LMStudioUnavailable (с упоминанием timeout)
|
||||
- Пустой стрим (только [DONE]): 0 delta'ов, 1 end
|
||||
- n>1 в стриме: WARNING в логах, payload уходит с n=1
|
||||
- n=0 валидация
|
||||
- Модель/usage/finish_reason приходят в end из последнего чанка
|
||||
- Skip не-SSE строк (комментариев/heartbeat)
|
||||
- Битый JSON в SSE → LMStudioUnavailable
|
||||
- reasoning-токены из content НЕ отделяются (проходят как обычный delta)
|
||||
|
||||
httpx мокается через unittest.mock: подменяем `httpx.Client`,
|
||||
а у экземпляра — `.stream(...)` (это context manager, отдаёт response
|
||||
с iter_lines).
|
||||
|
||||
Запуск: `python -m pytest tests/test_lm_streaming.py -v`
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
import pytest # noqa: E402
|
||||
|
||||
import httpx # noqa: E402
|
||||
|
||||
from lm_client import ( # noqa: E402
|
||||
LMStudioUnavailable,
|
||||
StreamEvent,
|
||||
stream_chat,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Хелперы для построения mock-стрима
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sse_chunk(*deltas: str, model: str = "qwen3.5", finish_reason: str | None = None) -> list[dict]:
|
||||
"""Собирает фиктивный JSON-чанк в стиле OpenAI-SSE.
|
||||
|
||||
Каждый delta — это content, который LM Studio кладёт в
|
||||
`choices[0].delta.content`. Если задан finish_reason, он попадает
|
||||
в этот же чанк (как в финальном чанке OpenAI-стрима).
|
||||
"""
|
||||
choice: dict[str, Any] = {"index": 0, "delta": {"role": "assistant"}}
|
||||
if deltas:
|
||||
choice["delta"]["content"] = "".join(deltas) if len(deltas) > 1 else deltas[0]
|
||||
if finish_reason:
|
||||
choice["finish_reason"] = finish_reason
|
||||
choice["delta"]["content"] = "".join(deltas) if deltas else ""
|
||||
return [{"id": "cmpl-x", "object": "chat.completion.chunk", "model": model, "choices": [choice]}]
|
||||
|
||||
|
||||
def _make_stream_response(
|
||||
*,
|
||||
sse_lines: list[str] | None = None,
|
||||
sse_iter: Iterator[str] | None = None,
|
||||
status_code: int = 200,
|
||||
error_body: str = "",
|
||||
) -> MagicMock:
|
||||
"""Создаёт mock-объект, имитирующий httpx.Response внутри `client.stream(...)`.
|
||||
|
||||
Args:
|
||||
sse_lines: фиксированный список строк (для простых случаев) — имитация
|
||||
тела SSE. Каждая строка — это ровно одна строка из `iter_lines()`.
|
||||
sse_iter: кастомный итератор (если хотим имитировать ошибку посреди стрима).
|
||||
status_code: HTTP status.
|
||||
error_body: тело для случая status_code >= 400 (читается через .read()).
|
||||
"""
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
|
||||
if sse_iter is not None:
|
||||
resp.iter_lines.return_value = sse_iter
|
||||
else:
|
||||
resp.iter_lines.return_value = sse_lines or []
|
||||
|
||||
if status_code >= 400:
|
||||
# resp.read() вызывается в stream_chat для превью ошибки.
|
||||
resp.read.return_value = error_body.encode("utf-8")
|
||||
else:
|
||||
resp.read.return_value = b""
|
||||
|
||||
return resp
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_stream(response: MagicMock | None = None, side_effect: Exception | None = None):
|
||||
"""Патчит `httpx.Client` так, что `client.stream(...)` отдаёт заданный response.
|
||||
|
||||
Args:
|
||||
response: mock-Response (то, что отдаёт __enter__ контекст-менеджера).
|
||||
side_effect: если задан — `client.stream(...)` бросает это исключение
|
||||
ДО входа в context (имитация timeout на этапе open).
|
||||
"""
|
||||
with patch("lm_client.httpx.Client") as MockClient:
|
||||
mock_inst = MagicMock()
|
||||
MockClient.return_value.__enter__.return_value = mock_inst
|
||||
|
||||
if side_effect is not None:
|
||||
mock_inst.stream.side_effect = side_effect
|
||||
else:
|
||||
# .stream(...) возвращает context manager, у которого __enter__
|
||||
# возвращает наш response.
|
||||
cm = MagicMock()
|
||||
cm.__enter__.return_value = response
|
||||
cm.__exit__.return_value = False
|
||||
mock_inst.stream.return_value = cm
|
||||
yield mock_inst
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Успешный стрим: 3 чанка + [DONE]
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_parses_three_sse_chunks():
|
||||
"""SSE из 3 чанков + [DONE] → 3 StreamEvent(type='delta') + 1 end."""
|
||||
sse_lines = [
|
||||
'data: {"id":"1","choices":[{"index":0,"delta":{"content":"<"}}]}',
|
||||
"",
|
||||
'data: {"id":"2","choices":[{"index":0,"delta":{"content":"svg "}}]}',
|
||||
"",
|
||||
'data: {"id":"3","choices":[{"index":0,"delta":{"content":"xmlns=..."}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
events = list(stream_chat(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
# 3 delta + 1 end = 4 события.
|
||||
assert len(events) == 4
|
||||
assert all(e.type == "delta" for e in events[:3])
|
||||
assert events[3].type == "end"
|
||||
|
||||
# Содержимое каждого delta.
|
||||
assert events[0].content == "<"
|
||||
assert events[1].content == "svg "
|
||||
assert events[2].content == "xmlns=..."
|
||||
|
||||
# Конкатенация = полный текст.
|
||||
full = "".join(e.content for e in events if e.type == "delta")
|
||||
assert full == "<svg xmlns=..."
|
||||
|
||||
# End-event: пустой content, но type='end'.
|
||||
assert events[3].content == ""
|
||||
assert events[3].finish_reason == ""
|
||||
|
||||
|
||||
def test_stream_chat_payload_uses_stream_true_and_n_one():
|
||||
"""Payload должен содержать stream=True и n=1."""
|
||||
sse_lines = ['data: {"choices":[{"delta":{"content":"x"}}]}', "", "data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
n=1,
|
||||
temperature=0.7,
|
||||
max_tokens=1024,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
call = mock_inst.stream.call_args
|
||||
body = call.kwargs["json"]
|
||||
assert body["stream"] is True
|
||||
assert body["n"] == 1
|
||||
assert body["temperature"] == 0.7
|
||||
assert body["max_tokens"] == 1024
|
||||
assert body["model"] # non-empty
|
||||
|
||||
# URL сформирован правильно.
|
||||
method, url = call.args[0], call.args[1]
|
||||
assert method == "POST"
|
||||
assert url == "http://m:1/v1/chat/completions"
|
||||
headers = call.kwargs["headers"]
|
||||
assert headers["Authorization"] == "Bearer lm-studio"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
|
||||
def test_stream_chat_collects_model_and_usage_in_end_event():
|
||||
"""Последний чанк содержит model/usage/finish_reason — это попадает в end."""
|
||||
sse_lines = [
|
||||
'data: {"model":"qwen3.5","choices":[{"delta":{"content":"Hel"}}]}',
|
||||
"",
|
||||
'data: {"model":"qwen3.5","choices":[{"delta":{"content":"lo"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}',
|
||||
"",
|
||||
'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert len(events) == 3 # 2 delta + 1 end
|
||||
end = events[-1]
|
||||
assert end.type == "end"
|
||||
assert end.model == "qwen3.5"
|
||||
assert end.usage == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
|
||||
assert end.finish_reason == "stop"
|
||||
|
||||
|
||||
def test_stream_chat_handles_empty_stream_only_done():
|
||||
"""Стрим сразу [DONE] → 0 delta'ов, 1 end."""
|
||||
sse_lines = ["data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "end"
|
||||
assert events[0].content == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Ошибки HTTP / network / timeout
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_500_raises_lmstudio_unavailable():
|
||||
"""HTTP 500 → LMStudioUnavailable, в сообщении есть 500 и body preview."""
|
||||
resp = _make_stream_response(status_code=500, error_body="internal error")
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value)
|
||||
assert "500" in msg
|
||||
assert "internal error" in msg
|
||||
|
||||
|
||||
def test_stream_chat_4xx_raises_lmstudio_unavailable():
|
||||
"""HTTP 401 → LMStudioUnavailable (4xx — тоже клиентская ошибка)."""
|
||||
resp = _make_stream_response(status_code=401, error_body="unauthorized")
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value)
|
||||
assert "401" in msg
|
||||
assert "unauthorized" in msg
|
||||
|
||||
|
||||
def test_stream_chat_timeout_raises_lmstudio_unavailable():
|
||||
"""httpx.TimeoutException на open стрима → LMStudioUnavailable с 'timeout'."""
|
||||
with _patched_stream(side_effect=httpx.TimeoutException("stream timed out")):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
base_url="http://m:1/v1",
|
||||
timeout_s=5.0,
|
||||
))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "timeout" in msg
|
||||
|
||||
|
||||
def test_stream_chat_network_error_raises_lmstudio_unavailable():
|
||||
"""Любой httpx.HTTPError → LMStudioUnavailable."""
|
||||
with _patched_stream(side_effect=httpx.ConnectError("connection refused")):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "недоступен" in msg or "unavailable" in msg
|
||||
assert "connection refused" in msg
|
||||
|
||||
|
||||
def test_stream_chat_broken_sse_json_raises_lmstudio_unavailable():
|
||||
"""Битый JSON в SSE-чанке → LMStudioUnavailable."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":"Hel"}}]}',
|
||||
"",
|
||||
'data: this is not json',
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
with pytest.raises(LMStudioUnavailable) as excinfo:
|
||||
list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
msg = str(excinfo.value).lower()
|
||||
assert "не-json" in msg or "sse" in msg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# n>1 в стриме → WARNING + n=1 в payload
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_n_greater_than_one_warns_and_uses_n_one(caplog):
|
||||
"""n>1 в стриме: логируется WARNING, payload уходит с n=1."""
|
||||
sse_lines = ['data: {"choices":[{"delta":{"content":"x"}}]}', "", "data: [DONE]", ""]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="lm_client"):
|
||||
with _patched_stream(response=resp) as mock_inst:
|
||||
events = list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
n=4,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
# Payload ушёл с n=1.
|
||||
body = mock_inst.stream.call_args.kwargs["json"]
|
||||
assert body["n"] == 1
|
||||
|
||||
# В логах было WARNING.
|
||||
assert any(
|
||||
"n=" in rec.message and "stream" in rec.message.lower()
|
||||
for rec in caplog.records if rec.levelno == logging.WARNING
|
||||
)
|
||||
|
||||
# Стрим всё равно отработал.
|
||||
assert any(e.type == "delta" for e in events)
|
||||
assert events[-1].type == "end"
|
||||
|
||||
|
||||
def test_stream_chat_rejects_n_less_than_one():
|
||||
"""n=0 → ValueError до обращения к сети."""
|
||||
with pytest.raises(ValueError, match="n должно быть"):
|
||||
list(stream_chat(
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
n=0,
|
||||
base_url="http://m:1/v1",
|
||||
))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Малые кейсы: пропуск мусорных строк, reasoning в content
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_chat_skips_non_data_lines():
|
||||
"""Строки без префикса `data:` (комментарии, event:, id:) — пропускаются."""
|
||||
sse_lines = [
|
||||
": this is a comment", # SSE-комментарий
|
||||
"event: message", # event-поле
|
||||
"id: 42", # id-поле
|
||||
'data: {"choices":[{"delta":{"content":"Hi"}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0].content == "Hi"
|
||||
assert events[-1].type == "end"
|
||||
|
||||
|
||||
def test_stream_chat_reasoning_tokens_pass_through_unchanged():
|
||||
"""qwen3.5 кладёт reasoning в основной content — мы НЕ отделяем."""
|
||||
reasoning_then_svg = (
|
||||
"Let me think about this. I need an icon of a fox. The viewBox is 64x64. "
|
||||
"<svg viewBox='0 0 64 64'><circle cx='32' cy='32' r='10'/></svg>"
|
||||
)
|
||||
# Один большой delta с reasoning+svg внутри.
|
||||
sse_lines = [
|
||||
f'data: {{"choices":[{{"delta":{{"content":{reasoning_then_svg!r}}}}}]}}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
full = "".join(e.content for e in events if e.type == "delta")
|
||||
assert "Let me think" in full
|
||||
assert "<svg" in full
|
||||
assert "</svg>" in full
|
||||
|
||||
|
||||
def test_stream_chat_handles_content_as_list_of_parts():
|
||||
"""content может быть list[dict] (мультимодальный стрим) — склеиваем в строку."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":[{"type":"text","text":"Hel"},{"type":"text","text":"lo"}]}}]}',
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert len(deltas) == 1
|
||||
assert deltas[0].content == "Hello"
|
||||
|
||||
|
||||
def test_stream_chat_emits_end_even_without_done_marker():
|
||||
"""Стрим без [DONE] всё равно получает end-event (с WARNING в логе)."""
|
||||
sse_lines = [
|
||||
'data: {"choices":[{"delta":{"content":"x"},"finish_reason":"stop"}]}',
|
||||
"",
|
||||
# нет [DONE] — обрыв по исчерпанию итератора
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
assert events[-1].type == "end"
|
||||
assert events[-1].finish_reason == "stop"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sanity: StreamEvent — frozen dataclass с нужными полями
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_stream_event_is_frozen_dataclass():
|
||||
"""StreamEvent — frozen: попытка изменения атрибута → FrozenInstanceError."""
|
||||
ev = StreamEvent(type="delta", content="x")
|
||||
assert ev.type == "delta"
|
||||
assert ev.content == "x"
|
||||
assert ev.usage is None
|
||||
assert ev.model == ""
|
||||
assert ev.finish_reason == ""
|
||||
with pytest.raises(Exception): # FrozenInstanceError
|
||||
ev.type = "end" # type: ignore[misc]
|
||||
|
||||
|
||||
def test_stream_chat_strip_data_prefix_with_or_without_space():
|
||||
"""SSE-префикс `data:` допускает опциональный пробел после двоеточия."""
|
||||
sse_lines = [
|
||||
'data:{"choices":[{"delta":{"content":"A"}}]}', # без пробела
|
||||
"",
|
||||
'data: {"choices":[{"delta":{"content":"B"}}]}', # с пробелом
|
||||
"",
|
||||
"data:[DONE]", # без пробела
|
||||
"",
|
||||
]
|
||||
resp = _make_stream_response(sse_lines=sse_lines)
|
||||
with _patched_stream(response=resp):
|
||||
events = list(stream_chat(messages=[{"role": "user", "content": "x"}], base_url="http://m:1/v1"))
|
||||
|
||||
deltas = [e for e in events if e.type == "delta"]
|
||||
assert [d.content for d in deltas] == ["A", "B"]
|
||||
assert events[-1].type == "end"
|
||||
Reference in New Issue
Block a user