diff --git a/py/services/agent/agent_service.py b/py/services/agent/agent_service.py index 853310ae..9b3a3313 100644 --- a/py/services/agent/agent_service.py +++ b/py/services/agent/agent_service.py @@ -284,6 +284,7 @@ class AgentService: model_path=model_path, llm_output=llm_response or {}, metadata=metadata, + readme_content=prompt_vars.get("readme_content_full", ""), ) if model_result.get("success", True): @@ -345,6 +346,7 @@ class AgentService: "hf_url": "", "repo": "", "readme_content": "", + "readme_content_full": "", "current_metadata": {}, "base_models": [], "priority_tags": "", @@ -367,6 +369,7 @@ class AgentService: if repo: readme = await self._fetch_readme(repo) context["readme_content"] = readme[:8000] if readme else "(README not available)" + context["readme_content_full"] = readme or "" try: context["base_models"] = await list_base_models() diff --git a/py/services/agent/post_processor.py b/py/services/agent/post_processor.py index be3e1593..e92f3ed2 100644 --- a/py/services/agent/post_processor.py +++ b/py/services/agent/post_processor.py @@ -39,15 +39,20 @@ class PostProcessor: model_path: str, llm_output: Dict[str, Any], metadata: Dict[str, Any], + readme_content: str = "", ) -> 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. + 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, + model_path, llm_output, metadata, readme_content, ) return { "success": False, @@ -64,12 +69,14 @@ class PostProcessor: model_path: str, llm_output: Dict[str, Any], metadata: Dict[str, Any], + readme_content: str = "", ) -> Dict[str, Any]: from ...agent_cli import ( apply_metadata_updates, download_preview, refresh_cache, ) + from .skills.enrich_hf_metadata.md_to_html import convert_readme_to_html updated_fields: List[str] = [] preview_downloaded = False @@ -97,12 +104,21 @@ class PostProcessor: civitai_updates["trainedWords"] = cleaned updates["civitai"] = civitai_updates - # modelDescription - new_desc = (llm_output.get("description") or "").strip() - if new_desc: - current_desc = metadata.get("modelDescription", "") or "" - if self._should_overwrite(current_desc, is_hf_model): - updates["modelDescription"] = new_desc + # modelDescription — from raw README content (converted to HTML) + if readme_content and is_hf_model: + converted = convert_readme_to_html(readme_content) + if converted: + updates["modelDescription"] = converted + + # short_description → civitai.description (for "About this version") + short_desc = (llm_output.get("short_description") or "").strip() + if short_desc and is_hf_model: + current_civitai = metadata.get("civitai") or {} + civitai_updates = dict(current_civitai) + if "civitai" in updates and isinstance(updates["civitai"], dict): + civitai_updates.update(updates["civitai"]) + civitai_updates["description"] = short_desc + updates["civitai"] = civitai_updates # tags new_tags = llm_output.get("tags", []) diff --git a/py/services/agent/skills/enrich_hf_metadata/SKILL.md b/py/services/agent/skills/enrich_hf_metadata/SKILL.md index 3666468d..9577d4f4 100644 --- a/py/services/agent/skills/enrich_hf_metadata/SKILL.md +++ b/py/services/agent/skills/enrich_hf_metadata/SKILL.md @@ -58,7 +58,7 @@ The trigger words or activation prompts needed to use this LoRA. Look for: - Example prompts at the start (usually the first word or phrase before any description) Return as an array of strings. If none found, return an empty array `[]`. **Never** return `["None"]` or any placeholder value — a truly empty list means no trigger words exist. -### description +### short_description A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. Return empty string if the README is too minimal. ### tags @@ -96,7 +96,7 @@ Return ONLY a JSON object with exactly these fields (no markdown fences, no extr "model_path": "{{model_path}}", "base_model": "", "trigger_words": ["", ""], - "description": "<1-2 sentence summary>", + "short_description": "<1-2 sentence summary>", "tags": ["", ""], "preview_url": "", "confidence": "" diff --git a/py/services/agent/skills/enrich_hf_metadata/md_to_html.py b/py/services/agent/skills/enrich_hf_metadata/md_to_html.py new file mode 100644 index 00000000..e3f5eeaa --- /dev/null +++ b/py/services/agent/skills/enrich_hf_metadata/md_to_html.py @@ -0,0 +1,300 @@ +"""Inline markdown-to-HTML converter for HF README content. + +No external dependencies. Strips YAML frontmatter, ```` sections, +badge images, and HTML comments before rendering. Only used by the +``enrich_hf_metadata`` skill. +""" + +from __future__ import annotations + +import html as html_module +import re +from typing import List, Tuple + + +def convert_readme_to_html(markdown_text: str | None) -> str: + """Convert HF README markdown to sanitised HTML.""" + if not markdown_text: + return "" + + text = markdown_text + text = _strip_frontmatter(text) + text = _strip_gallery(text) + text = _strip_badge_images(text) + text = _strip_html_comments(text) + html = _md_to_html(text) + return html.strip() + + +# --------------------------------------------------------------------------- +# Pre-processing: strip unwanted sections +# --------------------------------------------------------------------------- + + +def _strip_frontmatter(text: str) -> str: + if text.startswith("---"): + idx = text.find("---", 3) + if idx != -1: + return text[idx + 3 :] + return text + + +def _strip_gallery(text: str) -> str: + text = re.sub( + r"]*/>|]*>.*?", + "", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + text = re.sub( + r']*flex[^>]*>.*?(?:]*gallery).*?', + "", + text, + flags=re.IGNORECASE | re.DOTALL, + ) + return text + + +def _strip_badge_images(text: str) -> str: + badge_keywords = ( + "badge", "shield", "logo", "icon", "download", "license", + "python", "version", "status", "build", "test", "coverage", + "docker", "pypi", "npm", "github", "hugging face", "discord", + "twitter", "colab", "gradio", "space", + ) + + def _should_remove(m: re.Match) -> str: + alt = (m.group(1) or "").lower() + for kw in badge_keywords: + if kw in alt: + return "" + return m.group(0) + + text = re.sub(r"!\[([^\]]*)\]\([^)]+\)", _should_remove, text) + return text + + +def _strip_html_comments(text: str) -> str: + return re.sub(r"", "", text, flags=re.DOTALL) + + +# --------------------------------------------------------------------------- +# Markdown → HTML rendering +# --------------------------------------------------------------------------- + + +def _md_to_html(text: str) -> str: + """Convert a limited markdown subset to HTML. + + Handles: h1-h4, paragraphs, bold, italic, inline code, fenced code + blocks, unordered/ordered lists, links, images (→ alt text), hr, + blockquotes, and tables. + """ + lines = text.split("\n") + out: list[str] = [] + i = 0 + n = len(lines) + + while i < n: + line = lines[i] + + # Fenced code block + if line.startswith("```"): + lang = line[3:].strip() + code_lines: list[str] = [] + i += 1 + while i < n and not lines[i].startswith("```"): + code_lines.append(lines[i]) + i += 1 + i += 1 + code_escaped = html_module.escape("\n".join(code_lines)) + lang_attr = f' class="language-{html_module.escape(lang)}"' if lang else "" + out.append(f"
{code_escaped}\n
") + continue + + # Horizontal rule + if re.match(r"^-{3,}\s*$", line) or re.match(r"^\*{3,}\s*$", line): + out.append("
") + i += 1 + continue + + # ATX headers + h_match = re.match(r"^(#{1,4})\s+(.+?)\s*#*$", line) + if h_match: + level = len(h_match.group(1)) + content = _inline_md(h_match.group(2)) + out.append(f"{content}") + i += 1 + continue + + # Blockquote + if line.startswith("> "): + bq_lines: list[str] = [] + while i < n and lines[i].startswith("> "): + bq_lines.append(lines[i][2:]) + i += 1 + bq_html = _md_to_html("\n".join(bq_lines)) + out.append(f"
{bq_html}
") + continue + + # Unordered list + if re.match(r"^[-*+]\s", line): + items, i = _parse_list(lines, i, r"^[-*+]\s") + list_html = "".join(f"
  • {_inline_md(item)}
  • " for item in items) + out.append(f"
      {list_html}
    ") + continue + + # Ordered list + if re.match(r"^\d+\.\s", line): + items, i = _parse_list(lines, i, r"^\d+\.\s") + list_html = "".join(f"
  • {_inline_md(item)}
  • " for item in items) + out.append(f"
      {list_html}
    ") + continue + + # Table + if "|" in line and i + 1 < n and re.match(r"^\|[\s:-]+\|", lines[i + 1]): + table_html, i = _parse_table(lines, i) + out.append(table_html) + continue + + # Indented code block (4 spaces) — only when content is non-whitespace + if (line.startswith(" ") or line.startswith("\t")) and line.strip(): + code_lines, i = _collect_indented_code(lines, i) + if code_lines: + code_escaped = html_module.escape("\n".join(code_lines)) + out.append(f"
    {code_escaped}\n
    ") + continue + + # Whitespace-only indented line — treat as paragraph separator + if (line.startswith(" ") or line.startswith("\t")) and not line.strip(): + i += 1 + continue + + # Empty line + if not line.strip(): + i += 1 + continue + + # Paragraph (collect consecutive non-empty lines) + para_lines: list[str] = [] + while i < n: + l = lines[i] + if not l.strip(): + break + if re.match(r"^(#{1,4}\s|```|---|\*{3,}|\d+\.\s|[-*+]\s|>\s|\|)", l): + break + para_lines.append(l) + i += 1 + para = " ".join(p.strip() for p in para_lines if p.strip()) + if para: + out.append(f"

    {_inline_md(para)}

    ") + continue + + return "\n".join(out) + + +def _collect_indented_code(lines: list[str], start: int) -> Tuple[List[str], int]: + """Collect lines of an indented code block starting at *start*. + + Returns (code_lines, next_index). Whitespace-only lines within the + block are preserved; a block whose *only* content is whitespace + returns an empty list. + """ + code_lines: list[str] = [] + # Include the first line + if lines[start].strip(): + code_lines.append(lines[start]) + i = start + 1 + n = len(lines) + while i < n and (lines[i].startswith(" ") or lines[i].startswith("\t") or lines[i] == ""): + if lines[i].strip() or code_lines: + code_lines.append(lines[i]) + i += 1 + # If the block never had real content, discard it + if not any(ln.strip() for ln in code_lines): + return [], start + 1 + return code_lines, i + + +def _inline_md(text: str) -> str: + """Convert inline markdown within a paragraph/heading. + + Each pattern independently HTML-escapes its captured content, so no + global pre-escape is needed. This avoids double-escaping issues + (e.g. `` `a < b` `` → ``a < b``, not ``&lt;``). + """ + # Image: ![alt](url) → alt text + text = re.sub( + r"!\[([^\]]*)\]\([^)]+\)", + lambda m: html_module.escape(m.group(1) or ""), + text, + ) + + # Link: [text](url) → text + text = re.sub( + r"\[([^\]]+)\]\(([^)]+)\)", + lambda m: f'{html_module.escape(m.group(1))}', + text, + ) + + text = re.sub( + r"\*\*(.+?)\*\*", + lambda m: f"{html_module.escape(m.group(1))}", + text, + ) + text = re.sub( + r"(?{html_module.escape(m.group(1))}", + text, + ) + text = re.sub( + r"`([^`]+)`", + lambda m: f"{html_module.escape(m.group(1))}", + text, + ) + text = re.sub( + r"~~(.+?)~~", + lambda m: f"{html_module.escape(m.group(1))}", + text, + ) + return text + + +def _parse_list(lines: list[str], start: int, pattern: str) -> Tuple[List[str], int]: + """Collect list items starting at *start* matching *pattern*.""" + items: list[str] = [] + i = start + while i < len(lines): + m = re.match(pattern, lines[i]) + if m: + item = lines[i][m.end():].strip() + i += 1 + while i < len(lines) and lines[i].strip() and not re.match( + r"^(#{1,4}\s|```|\d+\.\s|[-*+]\s)", lines[i] + ): + item += " " + lines[i].strip() + i += 1 + items.append(item) + else: + break + return items, i + + +def _parse_table(lines: list[str], start: int) -> Tuple[str, int]: + """Parse a markdown table starting at *start*.""" + header_line = lines[start] + i = start + 2 # skip separator line + + headers = [c.strip() for c in header_line.strip().strip("|").split("|")] + rows: list[str] = [] + while i < len(lines) and "|" in lines[i]: + cells = [c.strip() for c in lines[i].strip().strip("|").split("|")] + cells_html = "".join(f"{_inline_md(c)}" for c in cells) + rows.append(f"{cells_html}") + i += 1 + + headers_html = "".join(f"{_inline_md(h)}" for h in headers) + table = f"{headers_html}" + table += "".join(rows) + table += "
    " + return table, i diff --git a/tests/agent_cli/test_agent_cli.py b/tests/agent_cli/test_agent_cli.py index 82b324db..a33f4470 100644 --- a/tests/agent_cli/test_agent_cli.py +++ b/tests/agent_cli/test_agent_cli.py @@ -315,3 +315,157 @@ class TestRefreshCache: mock_read.return_value = {} result = await refresh_cache("/some/path.safetensors") assert result is False + + +# ====================================================================== +# convert_readme_to_html — pure function, no mocks needed +# ====================================================================== + + +class TestConvertReadmeToHtml: + """Tests for the inline markdown→HTML converter.""" + + def test_empty_input(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + assert convert_readme_to_html("") == "" + assert convert_readme_to_html(None) == "" # type: ignore[arg-type] + + def test_heading(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + result = convert_readme_to_html("# Title") + assert "

    " in result and "Title" in result + + def test_subheadings(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "## Overview\n\n### Details" + result = convert_readme_to_html(md) + assert "

    Overview

    " in result + assert "

    Details

    " in result + + def test_bold_and_italic(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "**bold** and *italic*" + result = convert_readme_to_html(md) + assert "bold" in result + assert "italic" in result + + def test_inline_code(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "Use `model.train()`" + result = convert_readme_to_html(md) + assert "" in result and "model.train()" in result + + def test_fenced_code_block(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "```python\nprint('hello')\n```" + result = convert_readme_to_html(md) + assert "
    " in result
    +        assert "print" in result and "hello" in result
    +
    +    def test_unordered_list(self):
    +        from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
    +            convert_readme_to_html
    +
    +        md = "- item one\n- item two"
    +        result = convert_readme_to_html(md)
    +        assert "
      " in result + assert "
    • item one
    • " in result + assert "
    • item two
    • " in result + + def test_ordered_list(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "1. first\n2. second" + result = convert_readme_to_html(md) + assert "
        " in result + assert "
      1. first
      2. " in result + assert "
      3. second
      4. " in result + + def test_link(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "[click here](https://example.com)" + result = convert_readme_to_html(md) + assert 'click here' in result + + def test_badge_image_stripped(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "![badge](https://img.shields.io/badge/status-active)" + result = convert_readme_to_html(md) + assert "img.shields.io" not in result + + def test_gallery_stripped(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "Some text\n\nmore text" + result = convert_readme_to_html(md) + assert "Real content" in result + + def test_table(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "| A | B |\n|---|---|\n| 1 | 2 |" + result = convert_readme_to_html(md) + assert "" in result + assert "" in result + assert "" in result + + def test_horizontal_rule(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "before\n\n---\n\nafter" + result = convert_readme_to_html(md) + assert "
        " in result + + def test_inline_code_preserves_angle_bracket(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + result = convert_readme_to_html("Use `a < b` in code") + assert "a < b" in result + + def test_blockquote(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "> quoted text" + result = convert_readme_to_html(md) + assert "
        " in result + assert "quoted text" in result + + def test_indented_whitespace_not_treated_as_code(self): + from py.services.agent.skills.enrich_hf_metadata.md_to_html import \ + convert_readme_to_html + + md = "- item\n \n## heading after spacing" + result = convert_readme_to_html(md) + assert "
        " not in result
        +        assert "

        heading after spacing

        " in result diff --git a/tests/services/test_post_processor.py b/tests/services/test_post_processor.py index 54d34fa1..5f3437d5 100644 --- a/tests/services/test_post_processor.py +++ b/tests/services/test_post_processor.py @@ -67,7 +67,7 @@ class TestEnrichHfMetadata: MIN_LLM_OUTPUT = { "base_model": "", "trigger_words": [], - "description": "", + "short_description": "", "tags": [], "preview_url": "", "confidence": "low", @@ -167,23 +167,82 @@ class TestEnrichHfMetadata: applied = mock_apply.call_args[0][1] assert applied["civitai"]["trainedWords"] == ["trigger1", "trigger2"] - # -- description ----------------------------------------------------- + # -- short_description → civitai.description ------------------------- @pytest.mark.asyncio - async def test_description_set_when_empty(self, processor): - llm = {**self.MIN_LLM_OUTPUT, "description": "A model description"} + async def test_short_description_written_to_civitai(self, processor): + """short_description written to civitai.description for HF models.""" + llm = {**self.MIN_LLM_OUTPUT, "short_description": "A short summary"} 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( skill_name="enrich_hf_metadata", model_path="/p.safetensors", llm_output=llm, - metadata={"modelDescription": ""}, + metadata={"from_civitai": False}, ) - assert "modelDescription" in mock_apply.call_args[0][1] + applied = mock_apply.call_args[0][1] + assert applied["civitai"]["description"] == "A short summary" + + @pytest.mark.asyncio + async def test_short_description_skipped_for_civitai_model(self, processor): + """short_description NOT written for CivitAI models (has own description).""" + llm = {**self.MIN_LLM_OUTPUT, "short_description": "A short summary"} + with ( + mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply, + mock.patch("py.agent_cli.download_preview", return_value=None), + mock.patch("py.agent_cli.refresh_cache"), + ): + await processor.process( + skill_name="enrich_hf_metadata", + model_path="/p.safetensors", + llm_output=llm, + metadata={"from_civitai": True}, + ) + applied = mock_apply.call_args[0][1] + assert "civitai" not in applied or "description" not in applied.get("civitai", {}) + + # -- readme_content → modelDescription ------------------------------- + + @pytest.mark.asyncio + async def test_readme_content_converted_to_model_description(self, processor): + """Raw README converted to HTML and stored as modelDescription.""" + with ( + mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply, + mock.patch("py.agent_cli.download_preview", return_value=None), + mock.patch("py.agent_cli.refresh_cache"), + ): + await processor.process( + skill_name="enrich_hf_metadata", + model_path="/p.safetensors", + llm_output=self.MIN_LLM_OUTPUT, + metadata={"from_civitai": False}, + readme_content="# Hello\n\nThis is **bold**.", + ) + applied = mock_apply.call_args[0][1] + assert "

        Hello

        " in applied.get("modelDescription", "") + assert "bold" in applied.get("modelDescription", "") + + @pytest.mark.asyncio + async def test_readme_content_skipped_for_civitai_model(self, processor): + """README content NOT converted for CivitAI models.""" + with ( + mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply, + mock.patch("py.agent_cli.download_preview", return_value=None), + mock.patch("py.agent_cli.refresh_cache"), + ): + await processor.process( + skill_name="enrich_hf_metadata", + model_path="/p.safetensors", + llm_output=self.MIN_LLM_OUTPUT, + metadata={"from_civitai": True}, + readme_content="# Hello", + ) + applied = mock_apply.call_args[0][1] + assert "modelDescription" not in applied # -- tags ------------------------------------------------------------
        A1