Spaces:
Sleeping
Sleeping
| 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 RewriteResult, rewrite_memory | |
| from apps.p5_memory_quilt.quilt import render_tile | |
| def _dummy_card() -> RewriteResult: | |
| return RewriteResult( | |
| memory_text="Every Friday the tamale cart parked outside the blue house.", | |
| location_tag="corner store", | |
| style="Fabric Quilt", | |
| preferred_model_id="jetbrains/Mellm-2-12B-Instruct", | |
| fallback_model_id="openbmb/MiniCPM5-1B", | |
| selected_model_id="jetbrains/Mellm-2-12B-Instruct", | |
| checkpoint_path="/tmp/checkpoint", | |
| prompt_source="local-transformers", | |
| 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", | |
| keywords=("tamale", "house"), | |
| season_hint="soft evening light and a calm neighborhood mood", | |
| photo_path=None, | |
| ) | |
| def test_rewrite_memory_requires_a_local_checkpoint(monkeypatch: pytest.MonkeyPatch) -> None: | |
| monkeypatch.delenv("P5_MEMORY_QUILT_PRIMARY_MODEL_PATH", raising=False) | |
| monkeypatch.delenv("P5_MEMORY_QUILT_FALLBACK_MODEL_PATH", raising=False) | |
| with pytest.raises(RuntimeError, match="local checkpoint"): | |
| rewrite_memory("Every Friday the tamale cart parked outside the blue house.", "corner store", "Fabric Quilt") | |
| def test_rewrite_memory_uses_llm_backend_when_checkpoint_exists(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, | |
| "input_tokens": 11, | |
| "output_tokens": 22, | |
| }, | |
| ) | |
| 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") | |
| def test_render_tile_requires_a_local_flux_checkpoint(monkeypatch: pytest.MonkeyPatch) -> None: | |
| monkeypatch.delenv("P5_FLUX_MODEL_DIR", raising=False) | |
| with pytest.raises(RuntimeError, match="P5_FLUX_MODEL_DIR"): | |
| render_tile(_dummy_card()) | |