alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-07-06 09:21:16 -03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): rename md_to_html to readme_processor, fix section extraction, widget parsing, and list_base_models

Browse Source 
- Rename md_to_html.py → readme_processor.py (file no longer just HTML conversion)
- _extract_section: include YAML frontmatter, use heading-level-aware forward
  walk (sub-headings under # are included), increase walk limit past 30 lines
- _is_heading: exclude </hN> closing tags from boundary detection
- _heading_level: new helper for heading-level-aware section matching
- css: yield 0 for heading like closing tags, was unexpectedly caught by _is_heading
- extract_gallery_images: fix YAML block scalar (text: >-) prompt extraction;
  use endswith instead of == to detect the block marker
- _strip_widget_section: add to clean_readme_for_llm (widget text is handled
  by post-processor, not needed in LLM prompt)
- _strip_standalone_images: keep markdown image URLs intact for LLM preview
  extraction (was stripping to alt text only)
- list_base_models: switch from scanner-cache aggregation to
  CivitaiBaseModelService.get_base_models() - always returns full list
- Ollama: add num_ctx=32768 to payload options so thinking models have room
  to both reason and produce output
- Add tests/agent_cli/test_readme_processor.py: 59 tests covering extraction,
  cleaning, section matching, heading detection
- Update existing tests for behavioral changes
This commit is contained in:
Will Miao
2026-07-05 06:39:54 +08:00
parent 905c37290f
commit dd3aa97d0a
8 changed files with 733 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
py/agent_cli/__init__.py
View File

@@ -113,38 +113,29 @@ async def identify_model_type(model_path: str) -> str:
async def list_base_models(limit: int = 0) -> List[str]:
"""Return deduplicated base model names from all model caches.
"""Return all valid CivitAI base model names.
The result is ordered by frequency (most common first). Pass
*limit* = 0 (default) for all models.
Uses ``CivitaiBaseModelService.get_base_models()`` which merges a
hardcoded list (``SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS``) with remote
models fetched from the CivitAI API. Never empty — the hardcoded
fallback always provides a complete set.
The result is sorted alphabetically. Pass *limit* = 0 for all models.
"""
from ..services.service_registry import ServiceRegistry
from ..services.civitai_base_model_service import (
CivitaiBaseModelService,
)
counts: Dict[str, int] = {}
for getter_name in (
"get_lora_scanner",
"get_checkpoint_scanner",
"get_embedding_scanner",
):
getter = getattr(ServiceRegistry, getter_name, None)
if getter is None:
continue
try:
scanner = await getter()
if scanner is None:
continue
cache = await scanner.get_cached_data()
for entry in cache.raw_data:
bm = entry.get("base_model")
if bm:
counts[bm] = counts.get(bm, 0) + 1
except Exception as exc:
logger.debug("list_base_models scanner %s error: %s", getter_name, exc)
sorted_names = [name for name, _ in sorted(counts.items(), key=lambda x: -x[1])]
try:
service = await CivitaiBaseModelService.get_instance()
response = await service.get_base_models()
names: List[str] = response.get("models", [])
except Exception as exc:
logger.warning("list_base_models failed: %s", exc)
names = []
if limit > 0:
return sorted_names[:limit]
return sorted_names
return names[:limit]
return names
async def read_metadata(model_path: str) -> Dict[str, Any]:

6
py/services/agent/agent_service.py
View File

@@ -31,7 +31,7 @@ from ..llm_service import LLMService
from ..websocket_manager import ws_manager
from .post_processor import PostProcessor
from .skill_registry import SkillRegistry
from .skills.enrich_hf_metadata.md_to_html import (
from .skills.enrich_hf_metadata.readme_processor import (
clean_readme_for_llm,
extract_relevant_section,
)
@@ -397,6 +397,10 @@ class AgentService:
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 ""
logger.info(
"Cleaned README for %s (%d chars): ---BEGIN---\n%s\n---END---",
repo, len(cleaned), cleaned[:800] if cleaned else "(empty)",
)
try:
context["base_models"] = await list_base_models()

2
py/services/agent/post_processor.py
View File

@@ -78,7 +78,7 @@ class PostProcessor:
download_preview,
refresh_cache,
)
from .skills.enrich_hf_metadata.md_to_html import (
from .skills.enrich_hf_metadata.readme_processor import (
convert_readme_to_html,
extract_gallery_images,
extract_gallery_table_images,

99
py/services/agent/skills/enrich_hf_metadata/md_to_html.py → py/services/agent/skills/enrich_hf_metadata/readme_processor.py
View File

@@ -1,13 +1,8 @@
"""Inline markdown-to-HTML converter and LLM-prompt cleaner for HF README content.
"""HF README processing for the ``enrich_hf_metadata`` skill.
No external dependencies. Strips YAML frontmatter, ``<Gallery />`` sections,
badge images, and HTML comments before rendering. Used by the
``enrich_hf_metadata`` feature.
Also provides :func:`clean_readme_for_llm` which pre-processes the raw README
before it is injected into the LLM prompt, removing content that has zero value
for metadata extraction (widget sections, code blocks, training tables,
boilerplate, massive lists, etc.).
Provides README cleaning for LLM injection, gallery/image extraction from
multiple formats (YAML widget, markdown, HTML ``<img>``, gallery tables),
and section-based README trimming for collection repos.
"""
from __future__ import annotations
@@ -241,7 +236,26 @@ def extract_gallery_images(
if text_match:
raw_text = text_match.group(1).strip().strip("'\"")
if raw_text and raw_text != "-":
text = raw_text
# Handle YAML block scalar markers (>-, >, |, |-) where the
# actual text lives on subsequent indented lines.
if raw_text in (">", ">-", "|", "|-"):
text_lines: list[str] = []
in_block = False
for line in entry.split("\n"):
stripped = line.strip()
if not in_block:
if stripped.endswith(raw_text):
in_block = True
continue
# Block content ends at a line with less indentation
# or a YAML key at the start of a line.
if not stripped or re.match(r"^\s*\w+:", line):
break
if stripped:
text_lines.append(stripped)
text = " ".join(text_lines)
else:
text = raw_text
if url:
image: dict = {
@@ -439,6 +453,7 @@ def clean_readme_for_llm(markdown_text: str | None, max_length: int = 6000) -> s
# Order matters — broader strips first, then finer ones.
text = _strip_gallery(text)
text = _strip_widget_section(text)
text = _strip_fenced_code_blocks(text)
text = _strip_standalone_images(text)
text = _strip_training_tables(text)
@@ -722,6 +737,18 @@ def _looks_like_download_link(line: str) -> bool:
return False
def _heading_level(line: str) -> int:
"""Return the heading level of *line* (1-4), or 0 if not a heading."""
stripped = line.strip()
m = re.match(r"^(#{1,4})\s", stripped)
if m:
return len(m.group(1))
m = re.match(r"^<h([1-4])(?:\s|>)", stripped, re.IGNORECASE)
if m:
return int(m.group(1))
return 0
def _extract_section(
lines: list[str], match_idx: int, context_lines: int,
) -> str:
@@ -729,15 +756,23 @@ def _extract_section(
When *match_idx* is itself a heading line, the section starts *at*
that heading (no backward walk), avoiding pulling in content from
earlier sibling sections.
earlier sibling sections. The forward walk only stops at a heading
of **equal or higher** level (e.g. a ``#`` match includes all its
``##`` children).
Always includes the YAML frontmatter if the original lines contain one,
because it carries critical metadata (``base_model``, ``tags``,
``instance_prompt``) that the LLM needs regardless of which section
matches.
"""
n = len(lines)
# Determine start — if match is a heading, start right there
if _is_heading(lines[match_idx]):
start = match_idx
match_level = _heading_level(lines[match_idx])
else:
# Walk backward to find the nearest heading
match_level = 0
start = max(0, match_idx - context_lines)
for i in range(match_idx - 1, max(-1, match_idx - context_lines * 3), -1):
if i < 0:
@@ -747,13 +782,25 @@ def _extract_section(
start = i
break
# Walk forward to find the next heading at same or higher level
end = min(n, match_idx + context_lines)
for i in range(match_idx + 1, min(n, match_idx + context_lines * 3)):
if _is_heading(lines[i]):
# Walk forward. Stop at a heading of EQUAL or HIGHER (fewer #) level,
# so that a ``# Title`` match encompasses all its ``## Children``.
# Start from the full remaining lines so we don't truncate content
# when the YAML frontmatter pushes the matched heading far down.
end = n
walk_limit = min(n, match_idx + max(context_lines * 3, 120))
for i in range(match_idx + 1, walk_limit):
hl = _heading_level(lines[i])
if hl > 0 and (match_level == 0 or hl <= match_level):
end = i
break
# If YAML frontmatter exists before the matched section, prepend it.
if start > 0 and len(lines) > 1 and lines[0].strip() == "---":
for i in range(1, min(start, len(lines))):
if lines[i].strip() == "---":
yaml_section = "\n".join(lines[:i+1])
return yaml_section + "\n" + "\n".join(lines[start:end])
return "\n".join(lines[start:end])
@@ -801,6 +848,26 @@ def _strip_gallery(text: str) -> str:
return text
def _strip_widget_section(text: str) -> str:
"""Strip the ``widget:`` YAML block from the README frontmatter.
The widget section contains verbose example prompts (``text: >-`` entries)
that are useful for post-processor gallery image extraction but carry
no signal for LLM metadata extraction. Stripping them dramatically
reduces prompt size (e.g. 2800+ chars → ~100 chars) and lets the LLM
focus on the actual YAML metadata fields (``base_model``, ``tags``,
``instance_prompt``, etc.).
"""
# Match widget: through the end of the frontmatter (the closing ---)
# or until the next YAML top-level key.
return re.sub(
r"\nwidget:.*?(?=\n\w+:|\n---)",
"",
text,
flags=re.DOTALL,
)
def _strip_badge_images(text: str) -> str:
badge_keywords = (
"badge", "shield", "logo", "icon", "download", "license",

30
py/services/llm_service.py
View File

@@ -364,6 +364,9 @@ class LLMService:
"think": False,
"options": {
"temperature": temperature,
# Allow up to 32K context so the model has room to think
# AND produce output without hitting the 4K default limit.
"num_ctx": 32768,
},
}
if response_format is not None:
@@ -381,6 +384,16 @@ class LLMService:
if max_tokens is not None:
payload["max_tokens"] = max_tokens
if is_ollama:
logger.info(
"Ollama request: model=%s num_ctx=%s num_predict=%s format=%s think=%s",
payload.get("model"),
payload.get("options", {}).get("num_ctx"),
payload.get("options", {}).get("num_predict"),
payload.get("format", "none"),
payload.get("think"),
)
headers = self._build_headers(cfg["api_key"])
attempt = 0
@@ -507,8 +520,23 @@ class LLMService:
)
try:
return json.loads(result["content"])
parsed = json.loads(result["content"])
logger.info(
"LLM response base_model=%s tags=%s confidence=%s",
parsed.get("base_model", "?")[:50],
parsed.get("tags", []),
parsed.get("confidence", "?"),
)
logger.info(
"LLM raw content: %s",
(result.get("content") or "")[:1200],
)
return parsed
except (json.JSONDecodeError, TypeError) as exc:
logger.info(
"LLM raw response (first 800 chars): %s",
(result.get("content") or "")[:800],
)
logger.warning(
"LLM JSON parse failed on first attempt: %s. Retrying.", exc
)