Memory-Quilt / tests /test_trace_artifacts.py
Abhishek
Initialize project files and updated hackathon tags
485459a
Raw
History Blame Contribute Delete
5.24 kB
from __future__ import annotations
import importlib.util
import json
from pathlib import Path
from types import SimpleNamespace
import sys
ROOT = Path(__file__).resolve().parents[1]
for candidate in (ROOT, ROOT / "src"):
if str(candidate) not in sys.path:
sys.path.insert(0, str(candidate))
from apps.p5_memory_quilt.pipeline import QuiltState
from apps.p5_memory_quilt.prompt_rewriter import RewriteResult
def _load_headless_smoke_module():
script_path = ROOT / "scripts" / "headless_smoke.py"
spec = importlib.util.spec_from_file_location("p5_headless_smoke", script_path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Cannot import {script_path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _dummy_card(model_id: str, checkpoint_path: str, *, location_tag: str = "corner store") -> RewriteResult:
return RewriteResult(
memory_text="Every Friday the tamale cart parked outside the blue house.",
location_tag=location_tag,
style="Fabric Quilt",
preferred_model_id="jetbrains/Mellm-2-12B-Instruct",
fallback_model_id="openbmb/MiniCPM5-1B",
selected_model_id=model_id,
checkpoint_path=checkpoint_path,
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",
inference_meta={
"adapter_name": "local-transformers",
"backend": "local-transformers",
"model_id": model_id,
"checkpoint_path": checkpoint_path,
"checkpoint_source": "primary-checkpoint",
"generation_stats": {
"prompt_tokens": 11,
"generated_tokens": 22,
"elapsed_ms": 12.5,
},
},
photo_path=None,
)
def test_headless_smoke_writes_trace_artifact(tmp_path: Path, monkeypatch) -> None:
module = _load_headless_smoke_module()
module.TRACE_ARTIFACT_ROOT = tmp_path / "trace_artifacts"
pack_path = ROOT / "data" / "demo_packs" / "p5_memory_quilt"
output_dir = tmp_path / "verification"
checkpoint = tmp_path / "checkpoint"
checkpoint.mkdir()
seed_card = _dummy_card("jetbrains/Mellm-2-12B-Instruct", str(checkpoint))
smoke_card = _dummy_card("jetbrains/Mellm-2-12B-Instruct", str(checkpoint))
def fake_seed_state_from_pack(pack, count: int = 6):
return QuiltState(cards=(seed_card,), style=pack.style, pack_id=pack.pack_id)
def fake_add_memory_to_state(current_state, memory_text: str, location_tag: str = "", style: str | None = None, photo_path: str | None = None):
next_card = smoke_card
next_state = QuiltState(cards=current_state.cards + (next_card,), style=style or current_state.style, pack_id=current_state.pack_id)
return next_state, next_card
def fake_render_artifacts(state: QuiltState, output_dir: Path, stem: str = "memory_quilt"):
output_dir.mkdir(parents=True, exist_ok=True)
quilt_path = output_dir / f"{stem}_quilt.png"
tile_path = output_dir / f"{stem}_tile.png"
log_path = output_dir / f"{stem}_log.json"
quilt_path.write_bytes(b"quilt")
tile_path.write_bytes(b"tile")
log_path.write_text(json.dumps({"cards": len(state.cards)}), encoding="utf-8")
return SimpleNamespace(quilt_path=quilt_path, tile_path=tile_path, log_path=log_path, cards=state.cards)
monkeypatch.setattr(module, "seed_state_from_pack", fake_seed_state_from_pack)
monkeypatch.setattr(module, "add_memory_to_state", fake_add_memory_to_state)
monkeypatch.setattr(module, "render_artifacts", fake_render_artifacts)
payload = module.run_smoke(pack_path, output_dir, count=6)
trace_path = Path(payload["trace_path"])
assert payload["repo"] == "all4-p5-memory-quilt"
assert payload["project"] == "p5"
assert payload["kind"] == "smoke"
assert payload["network"] == "blocked"
assert payload["model_name"] == "jetbrains/Mellm-2-12B-Instruct"
assert payload["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
assert payload["adapter_name"] == "local-transformers"
assert payload["generation_stats"]["generated_tokens"] == 22
assert trace_path.exists()
trace = json.loads(trace_path.read_text(encoding="utf-8"))
assert trace["kind"] == "smoke"
assert trace["project"] == "p5"
assert trace["pack_id"] == "p5_memory_quilt"
assert trace["model_name"] == "jetbrains/Mellm-2-12B-Instruct"
assert trace["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
assert trace["adapter_name"] == "local-transformers"
assert trace["generation_stats"]["prompt_tokens"] == 11
assert trace["inputs"]["memory_count"] == trace["parsed_outputs"]["tile_count"]
assert Path(trace["parsed_outputs"]["quilt_path"]).exists()
assert Path(trace["parsed_outputs"]["tile_path"]).exists()
assert payload["trace_path"] == str(trace_path)