refactor(agent): apply model-source data without an LLM

`_build_prompt_context()` was only reached when the LLM was configured,
so a user with no provider got nothing at all from a linked model source
— no preview, no example images, no author summary, no tags — even
though all of that is deterministic data from a public API.

Split the model-card fetch into `_load_source_card()`, which runs for
every source-backed enrichment, and have the post-processor apply its
result whether or not the LLM runs. The prompt is then built from the
already-fetched card rather than re-fetching it.

Invoking "Enrich Metadata with AI" still always calls the provider; a
model source supplying a description, images and tags is not treated as
a reason to skip it, since the LLM's summary and notes are richer and an
action that silently does not call out to the provider would be
unpredictable. The site data acts as a fallback for the gaps the LLM
leaves.

Add `base_model_resolver.resolve_base_model()` to map the site's own
names (`krea/Krea-2-Turbo`, `KREA_2_TURBO`) onto the canonical
vocabulary, used only when the LLM returns no base model. It is strictly
conservative — exact normalised matching plus a bounded set of variant
suffixes, and it only ever returns a name that is already in the
vocabulary — so an uncertain hint defers to the LLM instead of writing a
plausible-looking wrong value.
This commit is contained in:
Will Miao
2026-09-14 20:39:10 +08:00
parent 35b291ab19
commit 4064ea7d3a
5 changed files with 756 additions and 23 deletions
+47 -4
View File
@@ -78,19 +78,62 @@ TensorArt is link-only: `tensor.art` sits behind a Cloudflare managed challenge
**What it does**:
1. Reads the model's `.metadata.json` to get the source (`source_platform` + `source_url`, or the legacy `hf_url`)
2. Fetches the model card through the provider in `py/services/model_sources/`
3. Sends the README + local metadata to the LLM for structured extraction
2. Fetches the model card through the provider in `py/services/model_sources/` — the README via `fetch_model_card()`, plus any extras the site keeps outside it via `fetch_model_card_context()`
3. Sends the README + site-provided extras + local metadata to the LLM for structured extraction
4. Writes extracted fields to `.metadata.json`:
- `base_model` — only if current value is empty
- `trainedWords` — trigger words (LoRA only, if none exist)
- `modelDescription`concise summary (if none exists)
- `modelDescription`the site's author description (if any) followed by the README rendered as HTML
- `tags` — merged with existing tags, deduplicated
- `civitai.images` — example images
- `metadata_source` — audit trail: `agent:enrich_hf_metadata`
- `llm_enriched_at` — ISO timestamp
5. Downloads and optimizes preview image (if LLM found one in the README)
5. Downloads and optimizes a preview image, using the per-file example image the
site publishes when the README has none
6. Updates the scanner cache
7. Broadcasts WebSocket progress events
#### Site-provided card extras (`fetch_model_card_context`)
A model card is not always just `README.md`. ModelScope keeps the author's
summary (`Description`), the site-curated tags (`OfficialTags`), and — per
published version — the model filenames together with that file's example
images (`MuseInfo.versions[].coverImages`) and trigger words in its
model-detail API. AIGC repositories there often ship an auto-generated
boilerplate README and put everything useful in `Description`, so reading only
the README yields almost nothing.
Providers opt in by overriding `ModelSource.fetch_model_card_context()`, which
returns a `ModelCardContext`. Example images are matched to the model's
**basename**, so each checkpoint in a collection repo gets its own images.
Sites with no such extras inherit an empty context, and the pipeline behaves
exactly as before.
#### Deterministic data is applied whether or not an LLM is configured
`AgentService._load_source_card()` runs for every source-backed enrichment, and
the post-processor applies what it returns before the LLM output is merged. A
user with **no** provider configured therefore still gets the author summary,
the example images, the preview, the site-curated tags, the trigger words and
the README rendered as the model description.
The LLM is always consulted when one is configured — invoking **Enrich Metadata
with AI** must call the provider every time, and the site data is never treated
as a reason to skip it. The deterministic values act as fallbacks that fill
gaps the LLM leaves behind:
| Field | Deterministic source | LLM role |
| --- | --- | --- |
| `modelDescription` | author summary + README as HTML | — |
| `civitai.images` | site example images, then README images | — |
| `preview_url` | first available example image | may propose one from the README |
| `tags` | site-curated tags, always merged in | proposes additional content tags |
| `civitai.description` | author summary | richer 1-2 sentence summary wins |
| `base_model` | site hints resolved against the canonical vocabulary (`py/services/agent/base_model_resolver.py`) | mapping it is the LLM's job; the resolver only fills in when the LLM returns nothing |
| `trainedWords` | per-file site trigger words, then YAML `instance_prompt` | primary extraction |
| `usage_tips` | regex over an explicitly stated strength range | primary extraction |
| `notes` | — | LLM-only |
Models with no source, an unknown source, or a source without model-card access (TensorArt) are skipped with an explicit reason and counted in the run summary.
**Model types**: LoRA, Checkpoint, Embedding
+134 -18
View File
@@ -27,11 +27,13 @@ from typing import Any, Dict, List, Optional
from ...config import config
from ..llm_service import LLMService
from ..model_sources import (
ModelCardContext,
get_source,
resolve_source_ref,
source_label,
)
from ..websocket_manager import ws_manager
from .base_model_resolver import resolve_base_model
from .post_processor import PostProcessor
from .skill_registry import SkillRegistry
from .skills.enrich_hf_metadata.readme_processor import (
@@ -282,14 +284,37 @@ class AgentService:
skip_model = True
if not skip_model:
prompt_vars: Dict[str, Any] = {"model_path": model_path}
if skill.llm_required and llm_configured:
prompt_vars = await self._build_prompt_context(
skill_name, model_path, metadata, registry, llm,
# The site's own data is deterministic and must land whether
# or not an LLM is available: a user without a key still gets
# the author summary, the example images and the tags.
source_vars, source_context = await self._load_source_card(
model_path, metadata,
)
resolved_base_model = ""
if skill_name == "enrich_hf_metadata" and not (
metadata.get("base_model") or ""
).strip():
resolved_base_model = await self._resolve_site_base_model(
source_context,
)
llm_response: Optional[Dict[str, Any]] = None
if skill.llm_required and llm_configured:
if skill.llm_required and not llm_configured:
# Without a provider the deterministic model-source data
# still lands; the LLM-only fields simply stay untouched.
logger.info(
"[%s] No LLM configured for %s — applying %s data only",
skill_name, model_filename,
"model-source"
if not source_context.is_empty()
else "README",
)
elif skill.llm_required:
prompt_vars = await self._build_prompt_context(
skill_name, model_path, metadata, registry, llm,
source_vars=source_vars,
source_context=source_context,
)
prompt_template = registry.load_prompt(skill_name)
rendered = _render_prompt(prompt_template, prompt_vars)
llm_response = await llm.chat_completion_json(
@@ -312,7 +337,9 @@ class AgentService:
model_path=model_path,
llm_output=llm_response or {},
metadata=metadata,
readme_content=prompt_vars.get("readme_content_full", ""),
readme_content=source_vars.get("readme_content_full", ""),
source_context=source_context,
resolved_base_model=resolved_base_model,
)
if model_result.get("success", True):
@@ -395,6 +422,77 @@ class AgentService:
"""
return "\n".join(f"- {m}" for m in models)
async def _load_source_card(
self, model_path: str, metadata: Dict[str, Any]
) -> tuple[Dict[str, Any], ModelCardContext]:
"""Fetch the model card and site-published extras for one model.
Runs for every source-backed enrichment regardless of LLM
availability, because everything it returns is deterministic data that
should be applied even without a configured provider.
"""
variables: Dict[str, Any] = {
"asset_base_url": "",
"source_description": "",
"source_base_model": "",
"source_official_tags": "",
"source_example_images": "",
"source_trigger_words": "",
"readme_content": "(README not available)",
"readme_content_full": "",
}
ref = resolve_source_ref(metadata)
source = get_source(ref.platform) if ref is not None else None
if ref is None or source is None or not source.supports_enrichment:
return variables, ModelCardContext()
raw_basename = os.path.splitext(os.path.basename(model_path))[0]
variables["asset_base_url"] = source.asset_base_url(ref.source_id)
readme = await source.fetch_model_card(ref.source_id)
# Sites such as ModelScope keep part of the model card outside the
# README (author summary, curated tags, per-file example images).
card_context = await source.fetch_model_card_context(
ref.source_id, os.path.basename(model_path)
)
variables["source_description"] = card_context.description
variables["source_base_model"] = card_context.base_model
variables["source_official_tags"] = "\n".join(
f"- {tag}" for tag in card_context.official_tags
)
variables["source_example_images"] = "\n".join(
f"- {url}" for url in card_context.example_images
)
variables["source_trigger_words"] = ", ".join(card_context.trigger_words)
# Trim README to the section relevant to this model file
# (collection repos often have multiple models in one README).
if readme and raw_basename:
trimmed = extract_relevant_section(readme, raw_basename)
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
else:
cleaned = clean_readme_for_llm(readme) if readme else ""
variables["readme_content"] = cleaned if cleaned else "(README not available)"
variables["readme_content_full"] = readme or ""
return variables, card_context
async def _resolve_site_base_model(self, source_context: ModelCardContext) -> str:
"""Resolve the site's base-model hints to a canonical name, or ``""``."""
from ...metadata_ops import list_base_models
hints = [*source_context.base_model_aliases, source_context.base_model]
if not any(hints):
return ""
try:
known_names = await list_base_models()
except Exception as exc:
logger.debug("Failed to list base models for site resolution: %s", exc)
return ""
return resolve_base_model(hints, known_names)
async def _build_prompt_context(
self,
skill_name: str,
@@ -402,16 +500,25 @@ class AgentService:
metadata: Dict[str, Any],
registry: SkillRegistry,
llm: Any,
*,
source_vars: Optional[Dict[str, Any]] = None,
source_context: Optional[ModelCardContext] = None,
) -> Dict[str, Any]:
"""Gather variables for the skill's prompt template.
Reads metadata, fetches the HF README (if applicable), lists available
Reads metadata, fetches the model card (unless a pre-fetched
*source_vars* / *source_context* pair is supplied), lists available
base models, loads user priority tags, and returns a dict that maps to
``{{variable}}`` placeholders in ``prompt.md``.
"""
from ...metadata_ops import identify_model_type, list_base_models
from ..settings_manager import SettingsManager
if source_vars is None or source_context is None:
source_vars, source_context = await self._load_source_card(
model_path, metadata,
)
context: Dict[str, Any] = {
"model_path": model_path,
"model_basename": "",
@@ -421,6 +528,15 @@ class AgentService:
"source_platform": "",
"source_label": "",
"asset_base_url": "",
# Site-provided card extras (see ModelSource.fetch_model_card_context)
"source_description": "",
"source_base_model": "",
"source_official_tags": "",
"source_example_images": "",
"source_trigger_words": "",
# Carrier for the structured context handed to the post-processor;
# never rendered into the prompt.
"source_context": ModelCardContext(),
# Legacy Hugging Face aliases (kept so older prompt templates and
# third-party skills keep rendering)
"hf_url": "",
@@ -458,17 +574,17 @@ class AgentService:
source = get_source(ref.platform) if ref is not None else None
if ref is not None and source is not None and source.supports_enrichment:
context["asset_base_url"] = source.asset_base_url(ref.source_id)
readme = await source.fetch_model_card(ref.source_id)
# Trim README to the section relevant to this model file
# (collection repos often have multiple models in one README).
if readme and raw_basename:
trimmed = extract_relevant_section(readme, raw_basename)
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
else:
cleaned = clean_readme_for_llm(readme) if readme else ""
context["readme_content"] = cleaned if cleaned else "(README not available)"
context["readme_content_full"] = readme or ""
# Values fetched once by _load_source_card and shared with the
# post-processor, so the network is not hit twice per model.
context["asset_base_url"] = source_vars["asset_base_url"]
context["source_context"] = source_context
context["source_description"] = source_vars["source_description"]
context["source_base_model"] = source_vars["source_base_model"]
context["source_official_tags"] = source_vars["source_official_tags"]
context["source_example_images"] = source_vars["source_example_images"]
context["source_trigger_words"] = source_vars["source_trigger_words"]
context["readme_content"] = source_vars["readme_content"]
context["readme_content_full"] = source_vars["readme_content_full"]
try:
raw_models = await list_base_models()
+94
View File
@@ -0,0 +1,94 @@
"""Map a site-reported base model onto this system's canonical vocabulary.
Model sites name base models in their own terms: ModelScope publishes
``krea/Krea-2-Turbo`` and ``KREA_2_TURBO`` where this system expects the
canonical ``Krea 2``. Turning one into the other is normally the LLM's job;
this module resolves the cases that can be decided safely so the canonical
field is still populated when the LLM returns nothing usable for it.
The resolver is deliberately strict, because a wrong base model written with
apparent authority is worse than no value at all:
* it only ever returns a name that is already present in *known_names*;
* matching is on the normalised form (lowercased, non-alphanumerics removed),
so separators and casing are ignored but nothing is inferred;
* a bounded set of published variant suffixes may be stripped, and only when
the remainder still matches a known name exactly.
Anything it cannot decide returns ``""``, and the caller falls back to the LLM.
"""
from __future__ import annotations
import re
from typing import Iterable, Sequence
#: Variant suffixes sites append to a base-model *family* name. Stripping one
#: is only attempted when the remainder matches a known name exactly, so an
#: unrecognised suffix can never produce a bogus match.
_VARIANT_SUFFIXES: tuple[str, ...] = (
"turbo",
"schnell",
"lightning",
"dev",
"beta",
"alpha",
)
_NON_ALNUM = re.compile(r"[^a-z0-9]+")
def _normalize(value: str) -> str:
"""Return the comparison form of *value*.
Lowercases and drops every non-alphanumeric character, so ``KREA_2``,
``Krea 2``, ``krea-2`` and ``krea.2`` all collapse to ``krea2``.
"""
return _NON_ALNUM.sub("", (value or "").lower())
def resolve_base_model(
hints: Iterable[str], known_names: Sequence[str]
) -> str:
"""Return the canonical base model that *hints* refers to, or ``""``.
Args:
hints: Site-reported names, best first (e.g. an architecture enum
before a link-style repository id).
known_names: The canonical vocabulary; only these are ever returned.
Returns:
One of *known_names*, or ``""`` when nothing matches exactly.
"""
normalized: dict[str, str] = {}
for name in known_names:
key = _normalize(name)
if key and key not in normalized:
normalized[key] = name
if not normalized:
return ""
ordered = [hint for hint in hints if hint]
# 1. Exact normalised match — the unambiguous case.
for hint in ordered:
candidate = _normalize(hint)
if candidate in normalized:
return normalized[candidate]
# 2. Drop one published variant suffix and retry exactly.
for hint in ordered:
candidate = _normalize(hint)
for suffix in _VARIANT_SUFFIXES:
if not candidate.endswith(suffix) or candidate == suffix:
continue
stem = candidate[: -len(suffix)]
if stem in normalized:
return normalized[stem]
return ""
__all__ = ["resolve_base_model"]
+384 -1
View File
@@ -11,6 +11,7 @@ from unittest import mock
import pytest
from py.services.agent.agent_service import AgentService
from py.services.model_sources import ModelCardContext
class TestEnrichmentSkipReason:
@@ -59,12 +60,23 @@ class TestBuildPromptContext:
async def test_modelscope_card_populates_source_variables(self):
service = AgentService()
readme = "---\nbase_model: krea/Krea-2-Turbo\n---\n# krea\n"
card_context = ModelCardContext(
description="权重0.5-1.2。配合《krea2-Cc-MJ-风格滤镜》lora一起使用。",
base_model="krea/Krea-2-Turbo",
official_tags=["photography", "woman"],
example_images=["https://resources.modelscope.cn/cover-images/a.png"],
trigger_words=["kreamodel"],
)
with (
mock.patch(
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
new=mock.AsyncMock(return_value=readme),
) as mock_fetch,
mock.patch(
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card_context",
new=mock.AsyncMock(return_value=card_context),
) as mock_context,
mock.patch(
"py.metadata_ops.list_base_models",
new=mock.AsyncMock(return_value=["Krea 2 Turbo"]),
@@ -91,6 +103,10 @@ class TestBuildPromptContext:
)
mock_fetch.assert_awaited_once_with("jj3550945163/Krea-2-LORA")
# The per-file lookup must receive the basename, not the full path.
mock_context.assert_awaited_once_with(
"jj3550945163/Krea-2-LORA", "krea.safetensors"
)
assert context["source_platform"] == "modelscope"
assert context["source_id"] == "jj3550945163/Krea-2-LORA"
assert context["source_label"] == "ModelScope"
@@ -99,10 +115,57 @@ class TestBuildPromptContext:
== "https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master"
)
assert readme in context["readme_content_full"]
# Site-provided extras are rendered into their own prompt variables.
assert context["source_description"] == card_context.description
assert context["source_base_model"] == "krea/Krea-2-Turbo"
assert context["source_official_tags"] == "- photography\n- woman"
assert (
context["source_example_images"]
== "- https://resources.modelscope.cn/cover-images/a.png"
)
assert context["source_trigger_words"] == "kreamodel"
# The structured context is carried through for the post-processor.
assert context["source_context"] is card_context
# Hugging Face aliases stay empty for a non-HF source.
assert context["hf_url"] == ""
assert context["repo"] == "jj3550945163/Krea-2-LORA"
@pytest.mark.asyncio
async def test_huggingface_card_context_is_empty(self):
"""Sources without card extras contribute empty prompt variables."""
service = AgentService()
with (
mock.patch(
"py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
new=mock.AsyncMock(return_value="# card\n"),
),
mock.patch(
"py.metadata_ops.list_base_models",
new=mock.AsyncMock(return_value=[]),
),
mock.patch(
"py.metadata_ops.identify_model_type",
new=mock.AsyncMock(return_value="lora"),
),
mock.patch(
"py.services.settings_manager.SettingsManager.get_priority_tag_config",
return_value={},
),
):
context = await service._build_prompt_context(
skill_name="enrich_hf_metadata",
model_path="/models/loras/thing.safetensors",
metadata={"hf_url": "https://huggingface.co/user/repo"},
registry=mock.Mock(),
llm=mock.Mock(),
)
assert context["source_description"] == ""
assert context["source_official_tags"] == ""
assert context["source_example_images"] == ""
assert context["source_context"].is_empty()
@pytest.mark.asyncio
async def test_huggingface_keeps_legacy_aliases(self):
service = AgentService()
@@ -179,4 +242,324 @@ class TestBuildPromptContext:
hf_fetch.assert_not_awaited()
ms_fetch.assert_not_awaited()
assert context["readme_content"] == ""
assert context["source_platform"] == "tensorart"
assert context["source_platform"] == "tensorart"
class TestExecuteSkillThreadsSourceContext:
"""The structured site context must reach the post-processor intact."""
@pytest.mark.asyncio
async def test_source_context_reaches_the_post_processor(self):
service = AgentService()
card_context = ModelCardContext(
example_images=["https://resources.modelscope.cn/cover-images/a.png"]
)
skill = mock.Mock(llm_required=True, input_schema={})
registry = mock.Mock()
registry.get_skill.return_value = skill
registry.load_prompt.return_value = "{{model_path}}"
llm = mock.Mock()
llm.is_configured.return_value = True
llm.chat_completion_json = mock.AsyncMock(return_value={"base_model": "Krea 2"})
source_vars = {
"readme_content_full": "# card",
"source_description": "",
"source_base_model": "",
"source_official_tags": "",
"source_example_images": "",
"source_trigger_words": "",
"asset_base_url": "",
"readme_content": "# card",
}
with (
mock.patch.object(
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
),
mock.patch.object(
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
),
mock.patch.object(
service,
"_load_source_card",
new=mock.AsyncMock(return_value=(source_vars, card_context)),
),
mock.patch.object(
service,
"_build_prompt_context",
new=mock.AsyncMock(
return_value={
"model_path": "/p.safetensors",
"readme_content_full": "# card",
"source_context": card_context,
}
),
),
mock.patch(
"py.metadata_ops.read_metadata",
new=mock.AsyncMock(
return_value={
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/u/r",
}
),
),
mock.patch(
"py.services.agent.agent_service.PostProcessor.process",
new=mock.AsyncMock(
return_value={"success": True, "updated_fields": []}
),
) as mock_process,
):
result = await service.execute_skill(
skill_name="enrich_hf_metadata",
input_data={"model_paths": ["/p.safetensors"]},
)
assert result.success is True
assert mock_process.call_args.kwargs["source_context"] is card_context
assert mock_process.call_args.kwargs["readme_content"] == "# card"
class TestSiteDataAppliedWithoutLlm:
"""A: the site's deterministic data must land even with no LLM available."""
@staticmethod
def _run(service, *, llm_configured: bool, card_context: ModelCardContext):
skill = mock.Mock(llm_required=True, input_schema={})
registry = mock.Mock()
registry.get_skill.return_value = skill
registry.load_prompt.return_value = "{{model_path}}"
llm = mock.Mock()
llm.is_configured.return_value = llm_configured
llm.chat_completion_json = mock.AsyncMock(return_value={"base_model": "Krea 2"})
source_vars = {
"readme_content_full": "# card",
"source_description": card_context.description,
"source_base_model": card_context.base_model,
"source_official_tags": "",
"source_example_images": "",
"source_trigger_words": "",
"asset_base_url": "",
"readme_content": "# card",
}
return skill, registry, llm, source_vars
@pytest.mark.asyncio
async def test_unconfigured_llm_still_applies_site_data(self):
service = AgentService()
card_context = ModelCardContext(
description="作者说明",
base_model="krea/Krea-2-Turbo",
official_tags=["photography"],
example_images=["https://cdn.example/a.png"],
)
skill, registry, llm, source_vars = self._run(
service, llm_configured=False, card_context=card_context
)
with (
mock.patch.object(
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
),
mock.patch.object(
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
),
mock.patch.object(
service,
"_load_source_card",
new=mock.AsyncMock(return_value=(source_vars, card_context)),
),
mock.patch.object(
service, "_resolve_site_base_model", new=mock.AsyncMock(return_value="Krea 2")
),
mock.patch(
"py.metadata_ops.read_metadata",
new=mock.AsyncMock(
return_value={
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/u/r",
}
),
),
mock.patch(
"py.services.agent.agent_service.PostProcessor.process",
new=mock.AsyncMock(
return_value={"success": True, "updated_fields": []}
),
) as mock_process,
):
result = await service.execute_skill(
skill_name="enrich_hf_metadata",
input_data={"model_paths": ["/p.safetensors"]},
)
assert result.success is True
# No LLM call, but the deterministic payload reached the post-processor.
llm.chat_completion_json.assert_not_awaited()
kwargs = mock_process.call_args.kwargs
assert kwargs["source_context"] is card_context
assert kwargs["readme_content"] == "# card"
assert kwargs["resolved_base_model"] == "Krea 2"
assert kwargs["llm_output"] == {}
class TestLoadSourceCard:
@pytest.mark.asyncio
async def test_returns_empty_variables_without_a_source(self):
service = AgentService()
variables, context = await service._load_source_card("/p.safetensors", {})
assert context.is_empty() is True
assert variables["readme_content_full"] == ""
assert variables["readme_content"] == "(README not available)"
@pytest.mark.asyncio
async def test_collects_readme_and_site_extras(self):
service = AgentService()
card = ModelCardContext(
description="作者说明",
base_model="krea/Krea-2-Turbo",
official_tags=["photography", "woman"],
example_images=["https://cdn.example/a.png"],
trigger_words=["kreaface"],
)
with (
mock.patch(
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
new=mock.AsyncMock(return_value="# card"),
),
mock.patch(
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card_context",
new=mock.AsyncMock(return_value=card),
) as mock_ctx,
):
variables, context = await service._load_source_card(
"/models/loras/krea.safetensors",
{
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/u/r",
},
)
mock_ctx.assert_awaited_once_with("u/r", "krea.safetensors")
assert context is card
assert variables["readme_content_full"] == "# card"
assert variables["source_description"] == "作者说明"
assert variables["source_official_tags"] == "- photography\n- woman"
assert variables["source_example_images"] == "- https://cdn.example/a.png"
assert variables["source_trigger_words"] == "kreaface"
assert variables["asset_base_url"].endswith("/resolve/master")
@pytest.mark.asyncio
async def test_resolve_site_base_model_uses_the_canonical_vocabulary(self):
service = AgentService()
with mock.patch(
"py.metadata_ops.list_base_models",
new=mock.AsyncMock(return_value=["Krea 2", "Flux.1 D"]),
):
resolved = await service._resolve_site_base_model(
ModelCardContext(
base_model="krea/Krea-2-Turbo",
base_model_aliases=["KREA_2", "KREA_2_TURBO"],
)
)
assert resolved == "Krea 2"
@pytest.mark.asyncio
async def test_resolve_site_base_model_is_empty_without_hints(self):
service = AgentService()
assert await service._resolve_site_base_model(ModelCardContext()) == ""
class TestLlmAlwaysRunsWhenConfigured:
"""Clicking "Enrich Metadata with AI" must always consult the LLM.
The site-provided data is applied deterministically either way, but it is
never treated as a reason to skip the call — the LLM's summary and notes
are richer than the raw site fields, and silently not calling out to the
provider would make the menu action unpredictable.
"""
@pytest.mark.asyncio
async def test_llm_runs_even_when_the_site_supplies_everything(self):
service = AgentService()
card_context = ModelCardContext(
description="作者说明",
base_model="krea/Krea-2-Turbo",
base_model_aliases=["KREA_2"],
official_tags=["photography", "woman"],
example_images=["https://cdn.example/a.png"],
trigger_words=["kreaface"],
)
skill = mock.Mock(llm_required=True, input_schema={})
registry = mock.Mock()
registry.get_skill.return_value = skill
registry.load_prompt.return_value = "{{model_path}}"
llm = mock.Mock()
llm.is_configured.return_value = True
llm.chat_completion_json = mock.AsyncMock(
return_value={"base_model": "Krea 2", "short_description": "llm summary"}
)
source_vars = {
"readme_content_full": "# card",
"source_description": card_context.description,
"source_base_model": card_context.base_model,
"source_official_tags": "- photography\n- woman",
"source_example_images": "- https://cdn.example/a.png",
"source_trigger_words": "kreaface",
"asset_base_url": "",
"readme_content": "# card",
}
with (
mock.patch.object(
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
),
mock.patch.object(
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
),
mock.patch.object(
service,
"_load_source_card",
new=mock.AsyncMock(return_value=(source_vars, card_context)),
),
mock.patch.object(
service,
"_build_prompt_context",
new=mock.AsyncMock(
return_value={"model_path": "/p.safetensors", "system_prompt": "sys"}
),
) as mock_prompt,
mock.patch(
"py.metadata_ops.read_metadata",
new=mock.AsyncMock(
return_value={
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/u/r",
"base_model": "Krea 2",
}
),
),
mock.patch(
"py.services.agent.agent_service.PostProcessor.process",
new=mock.AsyncMock(
return_value={"success": True, "updated_fields": []}
),
),
):
result = await service.execute_skill(
skill_name="enrich_hf_metadata",
input_data={"model_paths": ["/p.safetensors"]},
)
assert result.success is True
llm.chat_completion_json.assert_awaited_once()
# The prompt is built from the already-fetched card, not re-fetched.
assert mock_prompt.call_args.kwargs["source_context"] is card_context
assert mock_prompt.call_args.kwargs["source_vars"] is source_vars
@@ -0,0 +1,97 @@
"""Tests for the deterministic site base-model resolver.
The resolver exists so the enrichment pipeline can skip the LLM when the model
site already supplies everything; it must therefore be strictly conservative —
returning nothing is always better than returning the wrong canonical name.
"""
from __future__ import annotations
import pytest
from py.services.agent.base_model_resolver import resolve_base_model
KNOWN = [
"Krea 2",
"Flux.1 Krea",
"Flux.1 D",
"Flux.1 S",
"SDXL 1.0",
"Pony",
"Illustrious",
]
class TestExactNormalisedMatch:
@pytest.mark.parametrize(
"hint",
["Krea 2", "krea 2", "KREA_2", "krea-2", "krea.2", " Krea2 "],
)
def test_separators_and_casing_are_ignored(self, hint):
assert resolve_base_model([hint], KNOWN) == "Krea 2"
def test_returns_the_canonical_spelling_not_the_hint(self):
assert resolve_base_model(["kreA_2"], KNOWN) == "Krea 2"
def test_does_not_match_a_longer_prefixed_name_by_accident(self):
# "Flux.1 Krea" must not be resolved to "Krea 2".
assert resolve_base_model(["Flux.1 Krea"], KNOWN) == "Flux.1 Krea"
def test_first_matching_hint_wins(self):
assert (
resolve_base_model(["totally-unknown", "KREA_2"], KNOWN) == "Krea 2"
)
class TestVariantSuffixStripping:
@pytest.mark.parametrize(
"hint",
[
"KREA_2_TURBO",
"Krea-2-Turbo",
"krea2turbo",
"Krea 2 Turbo",
"krea-2-dev",
"krea2-schnell",
"krea2-lightning",
],
)
def test_common_published_suffixes_are_stripped(self, hint):
assert resolve_base_model([hint], KNOWN) == "Krea 2"
def test_suffix_only_hint_never_matches(self):
# "turbo" on its own strips to nothing and must not resolve.
assert resolve_base_model(["turbo"], KNOWN) == ""
class TestConservativeFailures:
@pytest.mark.parametrize(
"hints",
[
[],
[""],
["totally-unknown-model"],
["flux1dev"], # "Flux.1 D" normalises to "flux1d", not "flux1"
["sd"],
["ponyxl"],
],
)
def test_returns_empty_when_not_exactly_sure(self, hints):
assert resolve_base_model(hints, KNOWN) == ""
def test_returns_empty_without_a_vocabulary(self):
assert resolve_base_model(["KREA_2"], []) == ""
def test_only_ever_returns_a_known_name(self):
for name in ["KREA_2", "KREA_2_TURBO", "krea-2-turbo", "unknown"]:
result = resolve_base_model([name], KNOWN)
assert result == "" or result in KNOWN
def test_real_world_modelscope_hints(self):
"""The hints ModelScope actually publishes for a Krea 2 LoRA."""
assert (
resolve_base_model(
["KREA_2", "KREA_2_TURBO", "Krea-2-Turbo"], KNOWN
)
== "Krea 2"
)