feat(agent): improve enrich_hf_metadata skill with priority_tags, preview_url fix, civitai.trainedWords

- Add identify_model_type() helper to determine lora/checkpoint/embedding
- Pass priority_tags from user settings to LLM prompt for tag relevance
- SKILL.md: instruct LLM to exclude technical/generic HF tags, cross-reference
  against priority_tags; forbid ['None'] placeholder for trigger words
- post_processor: fix preview_url not updated after download (now writes local
  .webp path to metadata); write trigger words to civitai.trainedWords instead
  of top-level; sanitize ['None']/'null'/'n/a' placeholder values to []
- download_preview() now returns str | None (local path) instead of bool
- Update tests for new return type and nested civitai.trainedWords structure
This commit is contained in:
Will Miao
2026-07-02 22:14:44 +08:00
parent 63785f82b5
commit a8adcaf023
6 changed files with 121 additions and 44 deletions

View File

@@ -227,11 +227,11 @@ class TestApplyMetadataUpdates:
class TestDownloadPreview:
@pytest.mark.asyncio
async def test_empty_url_returns_false(self, tmp_path):
async def test_empty_url_returns_none(self, tmp_path):
mp = tmp_path / "m.safetensors"
mp.write_bytes(b"fake")
assert await download_preview(str(mp), "") is False
assert await download_preview(str(mp), " ") is False
assert await download_preview(str(mp), "") is None
assert await download_preview(str(mp), " ") is None
@pytest.mark.asyncio
async def test_successful_download_and_optimise(self, tmp_path):
@@ -246,12 +246,12 @@ class TestDownloadPreview:
get_dl.return_value = dl
exif.optimize_image.return_value = (b"optimized_webp", {})
result = await download_preview(str(mp), "https://ex.com/i.png")
assert result is True
assert result == str(tmp_path / "t.webp")
assert (tmp_path / "t.webp").exists()
assert (tmp_path / "t.webp").read_bytes() == b"optimized_webp"
@pytest.mark.asyncio
async def test_download_failure_returns_false(self, tmp_path):
async def test_download_failure_returns_none(self, tmp_path):
mp = tmp_path / "t.safetensors"
mp.write_bytes(b"fake")
with mock.patch("py.services.downloader.get_downloader") as get_dl:
@@ -260,7 +260,7 @@ class TestDownloadPreview:
dl.download_file = mock.AsyncMock(return_value=(False, None))
get_dl.return_value = dl
result = await download_preview(str(mp), "https://ex.com/i.png")
assert result is False
assert result is None
assert not (tmp_path / "t.webp").exists()

View File

@@ -44,7 +44,7 @@ class TestProcessDispatch:
mock.patch("py.agent_cli.refresh_cache") as mock_ref,
):
mock_apply.return_value = ["metadata_source"]
mock_dl.return_value = False
mock_dl.return_value = None
result = await processor.process(
skill_name="enrich_hf_metadata",
@@ -155,7 +155,7 @@ class TestEnrichHfMetadata:
llm = {**self.MIN_LLM_OUTPUT, "trigger_words": ["trigger1", "trigger2"]}
with (
mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply,
mock.patch("py.agent_cli.download_preview", return_value=False),
mock.patch("py.agent_cli.download_preview", return_value=None),
mock.patch("py.agent_cli.refresh_cache"),
):
await processor.process(
@@ -165,7 +165,7 @@ class TestEnrichHfMetadata:
metadata={"trainedWords": []},
)
applied = mock_apply.call_args[0][1]
assert applied["trainedWords"] == ["trigger1", "trigger2"]
assert applied["civitai"]["trainedWords"] == ["trigger1", "trigger2"]
# -- description -----------------------------------------------------
@@ -237,7 +237,7 @@ class TestEnrichHfMetadata:
mock.patch("py.agent_cli.download_preview") as mock_dl,
mock.patch("py.agent_cli.refresh_cache"),
):
mock_dl.return_value = True
mock_dl.return_value = "/p.webp"
result = await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
@@ -246,6 +246,8 @@ class TestEnrichHfMetadata:
)
assert result["preview_downloaded"] is True
mock_dl.assert_awaited_once_with("/p.safetensors", "https://ex.com/img.png")
applied = mock_apply.call_args[0][1]
assert applied["preview_url"] == "/p.webp"
@pytest.mark.asyncio
async def test_preview_skipped_when_exists(self, processor):