ComfyUI-Lora-Manager/py/services/agent/post_processor.py

"""Post-processing engine for skill pipeline outputs.

The :class:`PostProcessor` takes the LLM's structured JSON output and applies
it to a model's on-disk metadata via the :mod:`~py.metadata_ops` functions.

It handles all the skill-specific business logic — conditions, transformations,
and orchestration of multiple side-effects (write metadata, download preview,
refresh cache).  All actual I/O is delegated to :mod:`~py.metadata_ops`.
"""

from __future__ import annotations

import html
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:  # pragma: no cover - typing only
    from ..model_sources import ModelCardContext

logger = logging.getLogger(__name__)


class PostProcessor:
    """Deterministic post-processor for skill pipeline outputs.

    Usage (called by :class:`~py.services.agent.agent_service.AgentService`)::

        processor = PostProcessor()
        result = await processor.process(
            skill_name="enrich_hf_metadata",
            model_path="/path/to/model.safetensors",
            llm_output={...},
            metadata={...},    # from metadata_ops.read_metadata()
        )
    """

    async def process(
        self,
        *,
        skill_name: str,
        model_path: str,
        llm_output: Dict[str, Any],
        metadata: Dict[str, Any],
        readme_content: str = "",
        source_context: Optional["ModelCardContext"] = None,
        resolved_base_model: str = "",
        metadata_source: str = "agent:enrich_hf_metadata",
    ) -> Dict[str, Any]:
        """Route *llm_output* to the correct skill post-processor.

        *readme_content* is optional raw markdown content (e.g. HF README)
        that is converted to HTML and stored as ``modelDescription`` for
        the description tab.

        *source_context* carries the extras the model site publishes outside
        the README (author description, per-file example images, trigger
        words).  It is ``None`` for callers that have none.

        *resolved_base_model* is the canonical base-model name the site's own
        hints resolve to, used when the LLM did not supply one (which is the
        normal case when the LLM was skipped).

        *metadata_source* records who produced the metadata.  The AI skill
        keeps its historical value; the deterministic download-time hydration
        passes its own so the two remain distinguishable.  ``llm_enriched_at``
        is only stamped when *llm_output* actually carries a provider answer.

        Returns a dict with keys ``success`` (bool), ``updated_fields`` (list),
        ``preview_downloaded`` (bool), and ``errors`` (list).
        """
        if skill_name == "enrich_hf_metadata":
            return await self._process_enrich_hf_metadata(
                model_path, llm_output, metadata, readme_content, source_context,
                resolved_base_model, metadata_source,
            )
        return {
            "success": False,
            "updated_fields": [],
            "errors": [f"No post-processor registered for skill: {skill_name}"],
        }

    # ------------------------------------------------------------------
    # enrich_hf_metadata
    # ------------------------------------------------------------------

    async def _process_enrich_hf_metadata(
        self,
        model_path: str,
        llm_output: Dict[str, Any],
        metadata: Dict[str, Any],
        readme_content: str = "",
        source_context: Optional["ModelCardContext"] = None,
        resolved_base_model: str = "",
        metadata_source: str = "agent:enrich_hf_metadata",
    ) -> Dict[str, Any]:
        from ...metadata_ops import (
            apply_metadata_updates,
            download_preview,
            refresh_cache,
        )
        from ..model_sources import get_source, has_external_source, resolve_source_ref
        from .skills.enrich_hf_metadata.readme_processor import (
            convert_readme_to_html,
            extract_gallery_images,
            extract_gallery_table_images,
            extract_relevant_section,
            extract_simple_markdown_images,
            extract_html_img_tags,
        )

        updated_fields: List[str] = []
        preview_downloaded = False

        # -- Determine whether this is an externally-sourced model ---------
        # Key off the source fields directly: `from_civitai` records provenance
        # and can be true for a model that is also linked to an external site
        # (both sources coexist, see #1094), so it must not gate enrichment.
        is_source_model = has_external_source(metadata)

        source_ref = resolve_source_ref(metadata)
        source = get_source(source_ref.platform) if source_ref else None
        source_id = source_ref.source_id if source_ref else ""
        asset_base_url = (
            source.asset_base_url(source_id)
            if source is not None and source_id
            else None
        )

        # -- Collect updates -----------------------------------------------
        updates: Dict[str, Any] = {}

        # base_model — the LLM's mapping wins; when it returned nothing usable,
        # fall back to the canonical name the site's own hints resolve to.
        new_base = (llm_output.get("base_model") or "").strip()
        if not new_base:
            new_base = (resolved_base_model or "").strip()
        current_base = metadata.get("base_model", "") or ""
        if new_base and self._should_overwrite(current_base, is_source_model):
            updates["base_model"] = new_base

        # model_name — the site's own display name, so a source download never
        # shows up under its local filename.  Written only while the name is
        # still the untouched file stem: once a user renames a model that
        # choice is theirs to keep.
        site_name = ((source_context.model_name if source_context else "") or "").strip()
        if is_source_model and site_name:
            current_name = (metadata.get("model_name") or "").strip()
            file_stem = (metadata.get("file_name") or "").strip()
            if not current_name or current_name == file_stem:
                updates["model_name"] = site_name

        # trigger words → civitai.trainedWords
        new_triggers = llm_output.get("trigger_words", [])
        trigger_words_empty = True
        if isinstance(new_triggers, list):
            cleaned = [t.strip() for t in new_triggers if t.strip()]
            cleaned = [t for t in cleaned if t.lower() not in ("none", "null", "n/a")]
            trigger_words_empty = not cleaned
            current_triggers = (metadata.get("civitai") or {}).get("trainedWords") or []
            if self._should_overwrite_list(current_triggers, is_source_model):
                self._merge_civitai(updates, metadata, trainedWords=cleaned)

        # modelDescription — the author's own summary (when the site keeps one
        # outside the README, e.g. ModelScope's ``Description``) followed by the
        # README converted to HTML.
        site_description = (
            (source_context.description if source_context else "") or ""
        ).strip()
        if is_source_model and (site_description or readme_content):
            parts: List[str] = []
            if site_description:
                parts.append(f"<p>{html.escape(site_description)}</p>")
            if readme_content:
                converted = convert_readme_to_html(readme_content)
                if converted:
                    parts.append(converted)
            if parts:
                updates["modelDescription"] = "\n".join(parts)

        # short_description → civitai.description (for "About this version").
        # Falls back to the site's author summary, which for ModelScope AIGC
        # models is frequently the only human-written text available.
        short_desc = (llm_output.get("short_description") or "").strip()
        if not short_desc:
            short_desc = site_description
        if short_desc and is_source_model:
            self._merge_civitai(updates, metadata, description=short_desc)

        # The version label completes the card the way a CivitAI download does:
        # the UI renders `civitai.name` as the version chip.  It is per file,
        # so a collection repository shows that checkpoint's own label.
        site_version = (
            (source_context.version_name if source_context else "") or ""
        ).strip()
        if is_source_model and site_version:
            self._merge_civitai(updates, metadata, name=site_version)

        # gallery images → civitai.images (site example images, YAML frontmatter
        # widget entries, and Sample Gallery markdown tables in the README body)
        rec_width = llm_output.get("recommended_width") or 0
        rec_height = llm_output.get("recommended_height") or 0

        # Example images the site publishes for *this* file.  They are matched
        # by filename, so they are the most precise preview source available
        # and the only one for repositories whose README carries no images.
        site_images: List[Dict[str, Any]] = []
        if is_source_model and source_context is not None:
            site_images = [
                _example_image(url, rec_width, rec_height)
                for url in source_context.example_images
                if url
            ]

        gallery_images: List[Dict[str, Any]] = []
        if (readme_content or site_images) and is_source_model:
            repo = source_id
            readme_images: List[Dict[str, Any]] = []
            if readme_content and repo:
                # 1. Widget images (YAML frontmatter)
                gallery = extract_gallery_images(
                    readme_content, repo,
                    default_width=rec_width, default_height=rec_height,
                    base_url=asset_base_url,
                )

                # 2. Sample Gallery table images (markdown body), deduplicated
                existing_urls = {img["url"] for img in gallery if img.get("url")}
                table_images = extract_gallery_table_images(
                    readme_content, repo,
                    existing_urls=existing_urls,
                    default_width=rec_width, default_height=rec_height,
                    base_url=asset_base_url,
                )
                existing_urls.update(img["url"] for img in table_images if img.get("url"))

                # 3. Simple markdown images `![alt](url)` in the body
                simple_images = extract_simple_markdown_images(
                    readme_content, repo,
                    existing_urls=existing_urls,
                    default_width=rec_width, default_height=rec_height,
                    base_url=asset_base_url,
                )
                existing_urls.update(img["url"] for img in simple_images if img.get("url"))

                # 4. HTML `<img>` tags (used by many collection repos)
                html_images = extract_html_img_tags(
                    readme_content, repo,
                    existing_urls=existing_urls,
                    default_width=rec_width, default_height=rec_height,
                    base_url=asset_base_url,
                )

                readme_images = gallery + table_images + simple_images + html_images

            # Site images come first so the preview fallback below prefers an
            # image that is known to belong to this exact file.
            all_images = _dedupe_images(site_images + readme_images)
            if all_images:
                gallery_images = all_images
                self._merge_civitai(updates, metadata, images=all_images)

        # tags — the site's curated tags are authoritative content vocabulary, so
        # they are kept alongside whatever the LLM proposed (the LLM is skipped
        # entirely when the site data is complete, which is why this cannot rely
        # on ``llm_output`` alone).
        new_tags = llm_output.get("tags", [])
        candidate_tags: List[str] = []
        if is_source_model and source_context is not None:
            candidate_tags.extend(source_context.official_tags)
        if isinstance(new_tags, list):
            candidate_tags.extend(
                tag for tag in new_tags if tag not in candidate_tags
            )
        if candidate_tags:
            existing_tags = metadata.get("tags") or []
            merged = self._merge_tags(existing_tags, candidate_tags)
            if len(merged) > len(existing_tags) or is_source_model:
                updates["tags"] = merged

        # metadata_source is recorded for provenance; llm_enriched_at only means
        # something when a provider actually answered, so the deterministic
        # download-time hydration does not claim an enrichment that never ran.
        updates["metadata_source"] = metadata_source
        if llm_output:
            updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()

        # LLM confidence, stored for the enrichment evaluation harness.  The key
        # must NOT start with an underscore: `BaseModelMetadata.from_dict()`
        # deliberately drops underscore-prefixed keys so they never round-trip,
        # which silently erased this field on the next metadata write.
        raw_confidence = (llm_output.get("confidence") or "").strip()
        if raw_confidence:
            updates["llm_confidence"] = raw_confidence

        # Fallback: use the trigger words the site records for this exact file,
        # then the README's YAML `instance_prompt`, when the LLM returned none.
        if trigger_words_empty:
            site_triggers = (
                list(source_context.trigger_words) if source_context else []
            )
            if not site_triggers:
                instance_prompt = _extract_yaml_instance_prompt(readme_content)
                if instance_prompt:
                    site_triggers = [instance_prompt]
            if site_triggers:
                self._merge_civitai(updates, metadata, trainedWords=site_triggers)

        preview_remote_url = (llm_output.get("preview_url") or "").strip()
        # Fallback: if the LLM couldn't find a preview image in the cleaned
        # README, find the first gallery image from the *model-specific
        # section* of the README (not the repo-wide first image, which
        # belongs to a different model in collection repos).
        if not preview_remote_url and readme_content and is_source_model:
            model_basename = os.path.splitext(os.path.basename(model_path))[0]
            relevant_section = extract_relevant_section(
                readme_content, model_basename,
            )
            if relevant_section and relevant_section != readme_content:
                for img in gallery_images:
                    img_url = img.get("url", "")
                    if img_url and img_url in relevant_section:
                        preview_remote_url = img_url
                        break
        # Last resort: use the first gallery image from the full README.
        if not preview_remote_url and gallery_images:
            preview_remote_url = gallery_images[0].get("url", "")
        current_preview = metadata.get("preview_url") or ""
        if preview_remote_url and not (current_preview and os.path.exists(current_preview)):
            local_path = await download_preview(model_path, preview_remote_url)
            if local_path:
                preview_downloaded = True
                updates["preview_url"] = local_path

        # notes — plain-text summary of usage info from the LLM
        new_notes = (llm_output.get("notes") or "").strip()
        if new_notes:
            updates["notes"] = new_notes

        # usage_tips — JSON string (e.g. {"strength_min":0.85,"strength_max":1.4}).
        # When the LLM returned nothing, recover an explicitly stated strength
        # range from the author summary so the value is not lost.
        raw_tips = (llm_output.get("usage_tips") or "").strip()
        if not raw_tips or raw_tips == "{}":
            raw_tips = _extract_usage_tips(site_description)
        if raw_tips and raw_tips != "{}":
            try:
                json.loads(raw_tips)
                updates["usage_tips"] = raw_tips
            except (json.JSONDecodeError, TypeError):
                logger.warning(
                    "LLM returned invalid usage_tips JSON: %s", raw_tips[:200]
                )

        if updates:
            updated_fields = await apply_metadata_updates(model_path, updates)

        # -- Refresh scanner cache ------------------------------------------
        if updated_fields or preview_downloaded:
            await refresh_cache(model_path)

        return {
            "success": True,
            "updated_fields": updated_fields,
            "preview_downloaded": preview_downloaded,
            "updates": updates,
            "errors": [],
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _should_overwrite(current_value: str, is_source_model: bool) -> bool:
        """Return ``True`` when a scalar field should be overwritten."""
        return is_source_model or not current_value or current_value.lower() in (
            "", "unknown",
        )

    @staticmethod
    def _merge_civitai(
        updates: Dict[str, Any], metadata: Dict[str, Any], **fields: Any
    ) -> None:
        """Layer *fields* onto the ``civitai`` block being assembled.

        Description, version label, trigger words and gallery images all live
        in the same dict and are contributed by separate branches, so each one
        starts from what is already on disk and then applies whatever an
        earlier branch queued in *updates*.
        """

        merged = dict(metadata.get("civitai") or {})
        queued = updates.get("civitai")
        if isinstance(queued, dict):
            merged.update(queued)
        merged.update(fields)
        updates["civitai"] = merged

    @staticmethod
    def _should_overwrite_list(current_list: List[str], is_source_model: bool) -> bool:
        """Return ``True`` when a list field should be overwritten."""
        return is_source_model or not current_list

    @staticmethod
    def _merge_tags(existing: List[str], new: List[str]) -> List[str]:
        """Merge *new* tags into *existing*, all lowercased.

        This matches the behaviour of :class:`TagUpdateService` which
        normalises every tag to lowercase for case-insensitive dedup.
        """
        merged: List[str] = []
        seen: set[str] = set()
        for tag in list(existing) + list(new):
            t = tag.strip().lower()
            if t and t not in seen:
                merged.append(t)
                seen.add(t)
        return merged


# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------


#: Separator between a label and its value.  Published model cards routinely
#: wrap the numbers in markdown emphasis or quotes (``strength: **0.85 - 1.4**``,
#: ``CLIP 强度「0.5」``), so those are absorbed rather than treated as a break.
_EMPHASIS = "[\"'\u201c\u201d\u300c\u300d*_`\\s]*"

#: An explicitly stated strength/weight range, e.g. ``权重0.5-1.2``,
#: ``强度 0.8 ~ 1.2``, ``strength: **0.85 - 1.4**``.
_RANGE_DASH = "(?:-|\u2010|\u2011|\u2012|\u2013|\u2014|\uff0d|~|\uff5e|\u81f3|\u5230|to)"

_STRENGTH_RANGE_RE = re.compile(
    "(?:\u6743\u91cd|\u5f3a\u5ea6|strength|weight)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
    + r"(\d+(?:\.\d+)?)" + _EMPHASIS + _RANGE_DASH + _EMPHASIS
    + r"(\d+(?:\.\d+)?)",
    re.IGNORECASE,
)

#: A single strength/weight value, e.g. ``strength: 0.6``, ``权重 0.8``.
_STRENGTH_VALUE_RE = re.compile(
    "(?:\u6743\u91cd|\u5f3a\u5ea6|strength|weight)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
    + r"(\d+(?:\.\d+)?)",
    re.IGNORECASE,
)

#: ``clip strength: 0.5`` / ``CLIP 强度 0.5``.
_CLIP_STRENGTH_RE = re.compile(
    "clip" + _EMPHASIS + "(?:\u5f3a\u5ea6|strength)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
    + r"(\d+(?:\.\d+)?)",
    re.IGNORECASE,
)

#: ``clip skip: 2`` / ``CLIP 跳过 2``.
_CLIP_SKIP_RE = re.compile(
    "clip" + _EMPHASIS + "(?:skip|\u8df3\u8fc7)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
    + r"(\d+)",
    re.IGNORECASE,
)


def _extract_usage_tips(text: str) -> str:
    """Extract stated strength/CLIP recommendations from prose.

    This is the deterministic counterpart to the LLM's ``usage_tips`` output,
    used when the LLM was skipped.  It only recognises explicitly written
    values — it never infers a range — and returns ``""`` when it finds none.

    Returns:
        A JSON string matching the skill's ``usage_tips`` schema, or ``""``.
    """

    if not text:
        return ""

    tips: Dict[str, Any] = {}

    # CLIP strength is resolved first and then blanked out, so the generic
    # strength patterns cannot mistake `CLIP 强度 0.5` for the LoRA strength.
    text_for_strength = text
    clip_strength = _CLIP_STRENGTH_RE.search(text_for_strength)
    if clip_strength:
        tips["clip_strength"] = float(clip_strength.group(1))
        text_for_strength = (
            text_for_strength[: clip_strength.start()]
            + " "
            + text_for_strength[clip_strength.end() :]
        )

    range_match = _STRENGTH_RANGE_RE.search(text_for_strength)
    if range_match:
        low = float(range_match.group(1))
        high = float(range_match.group(2))
        if low > high:
            low, high = high, low
        tips["strength_min"] = low
        tips["strength_max"] = high
        tips["strength_range"] = f"{low:g}-{high:g}"
    else:
        value_match = _STRENGTH_VALUE_RE.search(text_for_strength)
        if value_match:
            tips["strength"] = float(value_match.group(1))

    clip_skip = _CLIP_SKIP_RE.search(text)
    if clip_skip:
        tips["clip_skip"] = int(clip_skip.group(1))

    if not tips:
        return ""
    return json.dumps(tips, ensure_ascii=False)


def _example_image(url: str, width: int, height: int) -> Dict[str, Any]:
    """Build a ``civitai.images`` entry for a site-provided example image.

    The site publishes no prompt alongside these images, so the entry carries
    empty prompt metadata and the LLM's recommended dimensions when it found
    any (falling back to the same 512px placeholder the README extractors use).
    """

    return {
        "url": url,
        "type": "image",
        "nsfwLevel": 0,
        "width": width or 512,
        "height": height or 512,
        "meta": {"prompt": "", "negativePrompt": ""},
        "hasMeta": False,
        "hasPositivePrompt": False,
    }


def _dedupe_images(images: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Drop later entries that repeat an earlier image URL, keeping order."""

    seen: set[str] = set()
    unique: List[Dict[str, Any]] = []
    for image in images:
        url = image.get("url") or ""
        if not url or url in seen:
            continue
        seen.add(url)
        unique.append(image)
    return unique


def _extract_yaml_instance_prompt(readme_content: str) -> str:
    """Extract ``instance_prompt`` from the YAML frontmatter of a HF README.

    Returns the prompt text, or empty string if not found.  Handles
    ``null`` / ``~`` YAML null values by returning empty string.
    """
    if not readme_content or not readme_content.startswith("---"):
        return ""

    # Find end of frontmatter
    end = readme_content.find("---", 3)
    if end == -1:
        return ""
    frontmatter = readme_content[3:end]

    for line in frontmatter.split("\n"):
        line = line.strip()
        m = re.match(r"^instance_prompt:\s*(.*)", line)
        if m:
            val = m.group(1).strip().strip('"').strip("'")
            if val.lower() in ("null", "~", "none", ""):
                return ""
            return val

    return ""