diff --git a/docs/agent_skills.md b/docs/agent_skills.md index 0ea5aab9..bcc264fd 100644 --- a/docs/agent_skills.md +++ b/docs/agent_skills.md @@ -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 diff --git a/py/services/agent/agent_service.py b/py/services/agent/agent_service.py index 05edc7a3..1440de0c 100644 --- a/py/services/agent/agent_service.py +++ b/py/services/agent/agent_service.py @@ -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() diff --git a/py/services/agent/base_model_resolver.py b/py/services/agent/base_model_resolver.py new file mode 100644 index 00000000..bf41da75 --- /dev/null +++ b/py/services/agent/base_model_resolver.py @@ -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"] diff --git a/tests/services/test_agent_enrichment_source.py b/tests/services/test_agent_enrichment_source.py index 0b6602ad..db8a9a93 100644 --- a/tests/services/test_agent_enrichment_source.py +++ b/tests/services/test_agent_enrichment_source.py @@ -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" \ No newline at end of file + 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 + diff --git a/tests/services/test_base_model_resolver.py b/tests/services/test_base_model_resolver.py new file mode 100644 index 00000000..fd2a75ec --- /dev/null +++ b/tests/services/test_base_model_resolver.py @@ -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" + )