Memory-Quilt / tests /test_prompt_rewriter.py
Abhishek
Initialize project files and updated hackathon tags
485459a
Raw
History Blame Contribute Delete
3.43 kB
from __future__ import annotations
from pathlib import Path
import sys
import pytest
ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
sys.path.insert(0, str(SRC))
from apps.p5_memory_quilt import llm_backend
from apps.p5_memory_quilt.prompt_rewriter import rewrite_memory
def test_rewrite_memory_uses_llm_backend_and_records_inference_meta(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
checkpoint = tmp_path / "checkpoint"
checkpoint.mkdir()
monkeypatch.setenv("P5_MEMORY_QUILT_PRIMARY_MODEL_PATH", str(checkpoint))
monkeypatch.delenv("P5_MEMORY_QUILT_FALLBACK_MODEL_PATH", raising=False)
captured: dict[str, object] = {}
def fake_try_rewrite(memory_text: str, location_tag: str, style: str, cfg: llm_backend.LLMConfig):
captured["memory_text"] = memory_text
captured["location_tag"] = location_tag
captured["style"] = style
captured["cfg"] = cfg
return (
{
"caption": "Corner Store",
"story": "A stitched memory of the corner store.",
"flux_prompt": "fabric quilt scene of a corner store at soft evening light",
"style_hint": "stitched fabric texture",
},
{
"adapter_name": "local-transformers",
"backend": "local-transformers",
"model_id": cfg.model_id,
"checkpoint_path": cfg.checkpoint_path,
"generation_stats": {
"prompt_tokens": 11,
"generated_tokens": 22,
"elapsed_ms": 12.5,
},
},
)
monkeypatch.setattr(llm_backend, "try_rewrite", fake_try_rewrite)
result = rewrite_memory("Every Friday the tamale cart parked outside the blue house.", "corner store", "Fabric Quilt")
assert captured["memory_text"] == "Every Friday the tamale cart parked outside the blue house."
assert captured["location_tag"] == "corner store"
assert captured["style"] == "Fabric Quilt"
assert isinstance(captured["cfg"], llm_backend.LLMConfig)
assert result.selected_model_id == "jetbrains/Mellm-2-12B-Instruct"
assert result.checkpoint_path == str(checkpoint)
assert result.prompt_source == "local-transformers"
assert result.caption == "Corner Store"
assert result.story == "A stitched memory of the corner store."
assert result.flux_prompt.startswith("fabric quilt scene")
assert result.inference_meta["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
assert result.inference_meta["checkpoint_source"] == "primary-checkpoint"
assert result.inference_meta["generation_stats"]["generated_tokens"] == 22
def test_rewrite_memory_rejects_empty_memory_text() -> None:
with pytest.raises(ValueError, match="memory_text"):
rewrite_memory("", "corner store", "Fabric Quilt")
def test_rewrite_memory_rejects_unknown_style(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
checkpoint = tmp_path / "checkpoint"
checkpoint.mkdir()
monkeypatch.setenv("P5_MEMORY_QUILT_PRIMARY_MODEL_PATH", str(checkpoint))
monkeypatch.delenv("P5_MEMORY_QUILT_FALLBACK_MODEL_PATH", raising=False)
with pytest.raises(ValueError, match="Unknown style"):
rewrite_memory("Every Friday the tamale cart parked outside the blue house.", "corner store", "Not A Style")