97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from config import Settings
|
|
from prompt_enhancer import PromptEnhancer
|
|
|
|
BASE = Path(__file__).resolve().parents[1] / "workflows"
|
|
|
|
|
|
class FakeEnhanceComfy:
|
|
def __init__(self, text: str = "A red apple on a wooden table, soft light") -> None:
|
|
self.text = text
|
|
self.workflow: dict | None = None
|
|
self.submits = 0
|
|
|
|
async def submit(self, workflow: dict, client_id: str) -> str:
|
|
self.workflow = workflow
|
|
self.submits += 1
|
|
return "prompt-1"
|
|
|
|
async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict:
|
|
return {"outputs": {"3": {"text": [self.text]}}}
|
|
|
|
def find_output_text(self, entry: dict) -> str | None:
|
|
return entry["outputs"]["3"]["text"][0]
|
|
|
|
|
|
def make_settings(enabled: bool = True) -> Settings:
|
|
return Settings(
|
|
bot_token="",
|
|
prompt_enhance_enabled=enabled,
|
|
workflow_prompt=BASE / "prompt_enhance.json",
|
|
)
|
|
|
|
|
|
async def test_enhance_t2i_prompt() -> None:
|
|
comfy = FakeEnhanceComfy()
|
|
enhancer = PromptEnhancer(comfy, make_settings())
|
|
|
|
result = await enhancer.enhance("красное яблоко на столе, мягкий свет", "t2i")
|
|
|
|
assert result == "A red apple on a wooden table, soft light"
|
|
assert comfy.submits == 1
|
|
assert comfy.workflow is not None
|
|
prompt = comfy.workflow["2"]["inputs"]["prompt"]
|
|
assert "<|im_start|>system" in prompt
|
|
assert "image generation model" in prompt
|
|
assert "красное яблоко" in prompt
|
|
assert prompt.endswith("<|im_start|>assistant\n")
|
|
|
|
|
|
async def test_enhance_i2i_prompt() -> None:
|
|
comfy = FakeEnhanceComfy("Make the apple blue, keep the composition")
|
|
enhancer = PromptEnhancer(comfy, make_settings())
|
|
|
|
result = await enhancer.enhance("Сделай яблоко синим", "i2i")
|
|
|
|
assert result == "Make the apple blue, keep the composition"
|
|
assert comfy.workflow is not None
|
|
prompt = comfy.workflow["2"]["inputs"]["prompt"]
|
|
assert "image-editing prompt translator" in prompt
|
|
|
|
|
|
async def test_enhance_with_context_uses_refine_prompt() -> None:
|
|
comfy = FakeEnhanceComfy("A blue apple on a wooden table, soft light")
|
|
enhancer = PromptEnhancer(comfy, make_settings())
|
|
|
|
result = await enhancer.enhance(
|
|
"Сделай яблоко синим",
|
|
"t2i",
|
|
previous_prompt="A red apple on a wooden table, soft light",
|
|
base_prompt="A red apple on a wooden table",
|
|
)
|
|
|
|
assert result == "A blue apple on a wooden table, soft light"
|
|
assert comfy.workflow is not None
|
|
prompt = comfy.workflow["2"]["inputs"]["prompt"]
|
|
assert "editing an existing image prompt" in prompt
|
|
assert "A red apple on a wooden table, soft light" in prompt
|
|
assert "Requested change: Сделай яблоко синим" in prompt
|
|
|
|
|
|
async def test_disabled_returns_original_without_comfy_call() -> None:
|
|
comfy = FakeEnhanceComfy()
|
|
enhancer = PromptEnhancer(comfy, make_settings(enabled=False))
|
|
|
|
result = await enhancer.enhance("красное яблоко", "t2i")
|
|
|
|
assert result == "красное яблоко"
|
|
assert comfy.submits == 0
|
|
|
|
|
|
async def test_clean_strips_prefix_and_empty_output() -> None:
|
|
assert PromptEnhancer._clean('Prompt: "A red apple"') == "A red apple"
|
|
assert PromptEnhancer._clean("user") == ""
|