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

feat(agent): enrich_hf_metadata — filename-aware section matching, preview extraction for markdown/HTML/widget, JSON salvage, instance_prompt fallback, and validation suite

Browse Source 
- extract_relevant_section(): trim README to model-filename-matching section
  for collection repos (download link, anchor ID, heading strategies)
- _strip_standalone_images(): preserve markdown image URLs so LLM can
  extract preview_url; strip only HTML <img> tags
- extract_simple_markdown_images(): extract civitai.images from ![]() body
- extract_html_img_tags(): extract from <img src="..."> (deadman44-style)
- extract_gallery_images(): fix widget parser for YAML - output: dash prefix
- _is_heading: exclude </hN> closing tags from boundary detection
- _extract_section: start at matching heading when match IS a heading line
- _try_salvage_json(): recover truncated JSON (close braces/brackets in
  LIFO order, close unterminated strings, strip trailing commas)
- PostProcessor: store _llm_confidence, add instance_prompt YAML fallback
- agent_service: pass model_basename to prompt, trim README via
  extract_relevant_section before clean_readme_for_llm
- Add tests/enrich_hf_validation/ suite: 100-model pipeline with progress
  checkpoint/resume, per-field scoring, markdown+JSON reporting
- Fix evaluation_engine: read _llm_confidence (not _llm_response)
This commit is contained in:
Will Miao
2026-07-04 12:00:15 +08:00
parent a1fd4e150b
commit 170c8068c5
13 changed files with 1998 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,11 +24,16 @@ from typing import Any, Dict, List, Optional
import aiohttp
import os
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 clean_readme_for_llm
from .skills.enrich_hf_metadata.md_to_html import (
clean_readme_for_llm,
extract_relevant_section,
)
logger = logging.getLogger(__name__)
@@ -344,6 +349,7 @@ class AgentService:
context: Dict[str, Any] = {
"model_path": model_path,
"model_basename": "",
"hf_url": "",
"repo": "",
"readme_content": "",
@@ -353,6 +359,11 @@ class AgentService:
"priority_tags": "",
}
# Extract model basename (filename without extension) for the LLM
# to use when locating the matching section in collection repos.
raw_basename = os.path.splitext(os.path.basename(model_path))[0]
context["model_basename"] = raw_basename or ""
context["current_metadata"] = {
"file_name": metadata.get("file_name", ""),
"base_model": metadata.get("base_model", ""),
@@ -369,7 +380,13 @@ class AgentService:
context["repo"] = repo or ""
if repo:
readme = await self._fetch_readme(repo)
cleaned = clean_readme_for_llm(readme) if readme else ""
# 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 ""

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

@@ -13,6 +13,7 @@ from __future__ import annotations
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
@@ -81,6 +82,8 @@ class PostProcessor:
convert_readme_to_html,
extract_gallery_images,
extract_gallery_table_images,
extract_simple_markdown_images,
extract_html_img_tags,
extract_repo_from_hf_url,
)
@@ -101,9 +104,11 @@ class PostProcessor:
# 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_civitai = metadata.get("civitai") or {}
current_triggers = current_civitai.get("trainedWords") or []
if self._should_overwrite_list(current_triggers, is_hf_model):
@@ -152,8 +157,24 @@ class PostProcessor:
existing_urls=existing_urls,
default_width=rec_w, default_height=rec_h,
)
existing_urls.update(img["url"] for img in table_images if img.get("url"))
all_images = gallery + table_images
# 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_w, default_height=rec_h,
)
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_w, default_height=rec_h,
)
all_images = gallery + table_images + simple_images + html_images
if all_images:
gallery_images = all_images
current_civitai = metadata.get("civitai") or {}
@@ -175,6 +196,23 @@ class PostProcessor:
updates["metadata_source"] = "agent:enrich_hf_metadata"
updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()
# Store LLM confidence in metadata so it's accessible for evaluation
raw_confidence = (llm_output.get("confidence") or "").strip()
if raw_confidence:
updates["_llm_confidence"] = raw_confidence
# Fallback: extract instance_prompt from YAML frontmatter when the LLM
# returned empty trigger words but the README has instance_prompt.
if trigger_words_empty:
instance_prompt = _extract_yaml_instance_prompt(readme_content)
if instance_prompt:
current_civitai = metadata.get("civitai") or {}
trig_civitai = dict(current_civitai)
if "civitai" in updates and isinstance(updates["civitai"], dict):
trig_civitai.update(updates["civitai"])
trig_civitai["trainedWords"] = [instance_prompt]
updates["civitai"] = trig_civitai
preview_remote_url = (llm_output.get("preview_url") or "").strip()
# Fallback: if the LLM couldn't find a preview image in the cleaned
# README, use the first gallery image extracted from the YAML widget
@@ -249,3 +287,35 @@ class PostProcessor:
merged.append(t)
seen.add(t)
return merged
# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------
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 ""

62
py/services/agent/skills/enrich_hf_metadata/SKILL.md
View File

@@ -13,6 +13,7 @@ You are an expert assistant for AI image generation models. Your task is to extr
- **Repository**: {{hf_url}}
- **Model file path**: {{model_path}}
- **Model filename**: {{model_basename}}
- **Repository ID**: {{repo}}
## Current Metadata (may be incomplete)
@@ -36,6 +37,31 @@ These are the subjects, styles, and concepts the user considers useful for categ
The following base models are currently valid in this system:
{{base_models}}
## HuggingFace → CivitAI Base Model Mapping
HuggingFace repos often declare `base_model:` in YAML frontmatter using HuggingFace
model names. Map them to CivitAI names using this reference:
| HuggingFace repo / model name | CivitAI base model |
|-------------------------------------------------------|--------------------|
| `runwayml/stable-diffusion-v1-5` | SD 1.5 |
| `CompVis/stable-diffusion-v1-4` | SD 1.5 |
| `stabilityai/stable-diffusion-2-1` / `-2-1-base` | SD 2.1 |
| `stabilityai/stable-diffusion-xl-base-1.0` | SDXL 1.0 |
| `stabilityai/stable-diffusion-3-5-large` | SD 3.5 Large |
| `black-forest-labs/FLUX.1-dev` | Flux.1 D |
| `black-forest-labs/FLUX.1-schnell` | Flux.1 S |
| `black-forest-labs/FLUX.1-krea` | Flux.1 Krea |
| `krea/Krea-2-Raw` or `krea/Krea-2-Turbo` | Krea 2 |
| `Comfy-Org/z_image_turbo` | ZImageTurbo |
| `alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1` | ZImageTurbo |
| `CogVideo` | CogVideoX |
| `Wan-AI/Wan2.1-T2V-14B` | Wan Video 2.2 T2V-A14B |
| `stabilityai/stable-video-diffusion-img2vid` | SVD |
The model file name itself may also hint at the base model (e.g. "flux", "sdxl",
"sd15", "krea2" in the filename).
## HuggingFace README Content
```
@@ -49,24 +75,26 @@ Extract the following information from the README content above:
### base_model
The base model this model was trained on. Use EXACTLY one of the names from the **Available Base Models** list above. Do not invent new names or use aliases.
Check the YAML frontmatter (between --- markers) for `base_model:` first, then look at the description text and safetensors metadata. If you cannot determine it, return an empty string.
Check the YAML frontmatter (between --- markers) for `base_model:` first, then look at the description text and safetensors metadata. If the YAML frontmatter uses a HuggingFace model name (e.g. `runwayml/stable-diffusion-v1-5`), use the mapping table above to find the correct CivitAI name. If you cannot determine it, return an empty string.
### trigger_words
The trigger words or activation prompts needed to use this LoRA. Look for:
- `instance_prompt:` in the YAML frontmatter
- Phrases like "trigger word:", "trigger:", "use this prompt:", "activation prompt:"
- In collection repos: the trigger section **specific to this model file** (look near matching download links or anchor IDs)
- 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.
### 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.
A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`, not the repo as a whole. Return empty string if the README is too minimal.
### tags
3-8 relevant tags for categorizing this model. **Quality over quantity.**
Sources to consider:
- The YAML frontmatter `tags:` list
- The YAML frontmatter `tags:` list (filter out technical ones — see below)
- The subject, style, character, or concept the model represents
- The model filename itself may give clues (e.g. "pokemon", "anime", "pixelart")
**Critical filtering rules — apply them strictly:**
@@ -82,10 +110,15 @@ Return empty array if no meaningful content tags remain after filtering.
The recommended image generation resolution for this model, in pixels. Look for sections like "Best Dimensions", "Recommended size", "Suggested resolution", or similar phrasing in the README. Prefer the explicitly marked "Best" or default resolution. If the table/list has multiple entries (e.g. "768 x 1024 (Best)" and "1024 x 1024 (Default)"), use the one marked "Best". Return integers. If no resolution can be determined, return 0 for both.
### preview_url
The URL of the most suitable preview image from the README. Look for image tags (e.g. `![alt](url)`) and the YAML frontmatter `widget:` section (which often has `output.url` fields). Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
The URL of the most suitable preview image from the README. Look for:
- Image tags near the section matching the model filename (`{{model_basename}}`)
- The YAML frontmatter `widget:` section (which often has `output.url` fields)
- In collection repos: the sample images listed **under the section** for this specific model version
- Generic `![alt](url)` in the body
Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
### notes
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. Return empty string if the README has no useful usage info.
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`. Return empty string if the README has no useful usage info.
### usage_tips
A JSON string with structured usage recommendations. Extract from the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5"). Possible fields (include only those you can determine):
@@ -109,6 +142,25 @@ Your confidence level in the extracted data:
- "medium" — some fields were inferred from context
- "low" — most fields are guesses based on limited information
## Important: Handling Collection Repos (multiple model files)
Many HuggingFace repos contain **multiple model files** in a single repository
(e.g. a "LoRA collection" with different styles/characters in separate files).
The model file currently being enriched is: **`{{model_basename}}`**
To find the correct section in the README:
1. **Search for download links** containing the filename — the surrounding paragraph is your section.
2. **Search for anchor IDs** (`<a id="...">`) or section headings whose text matches words from the filename.
3. **Search for HTML headings** (`<h1>`, `<h2>`, `<span>`) containing parts of the filename.
4. If no match is found, use the full README as usual — the model may be the only one in the repo.
When a matching section IS found, prefer metadata from that section.
When no section matches (e.g. single-model repos or repos without per-file sections),
extract metadata from the full README normally. Do not return empty data just
because the filename doesn't appear in the README.
## Output Format
Return ONLY a JSON object with exactly these fields (no markdown fences, no extra text):

368
py/services/agent/skills/enrich_hf_metadata/md_to_html.py
View File

@@ -20,6 +20,155 @@ from typing import List, Tuple
_REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
def extract_simple_markdown_images(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
"""Extract standalone markdown images from the README body.
Matches ``![alt](url)`` on lines that are NOT part of a markdown table
and NOT inside fenced code blocks. These are common in DreamBooth
training dumps where the user uploads example images with simple
``![img_0](./image_0.png)`` syntax.
Returns a list of dicts in the same ``civitai.images`` format as
:func:`extract_gallery_images`.
"""
if not markdown_text or not repo:
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
# Collect lines that are NOT inside fenced code blocks
lines = markdown_text.split("\n")
in_code_block = False
n = len(lines)
i = 0
while i < n:
line = lines[i]
if line.strip().startswith("```"):
in_code_block = not in_code_block
i += 1
continue
if in_code_block:
i += 1
continue
# Skip table rows
if "|" in line and i + 1 < n and re.match(r"^\|[\s:-]+\|", lines[i + 1]):
i += 1
continue
m = re.match(r"^\s*!\[([^\]]*)\]\(([^)]+)\)\s*$", line)
if m:
raw_path = m.group(2).strip()
if raw_path.startswith("http"):
url = raw_path
else:
clean = raw_path.lstrip("./").lstrip("/")
url = f"{base_url}/{clean}"
if url not in seen_urls:
seen_urls.add(url)
images.append({
"url": url,
"type": "image",
"nsfwLevel": 0,
"width": default_width,
"height": default_height,
"meta": {"prompt": m.group(1), "negativePrompt": ""},
"hasMeta": bool(m.group(1)),
"hasPositivePrompt": bool(m.group(1)),
})
i += 1
return images
def extract_html_img_tags(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
"""Extract image URLs from HTML ``<img src=\"...\">`` tags in the README.
Many HF collection repos (e.g. ``deadman44/Z-Image_LoRA``) use raw HTML
``<img>`` tags exclusively for their sample images, with no markdown
``![]()`` equivalents. This function finds those tags and constructs
resolvable HF URLs.
Returns a list of dicts in the ``civitai.images`` format.
"""
if not markdown_text or not repo:
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
for m in re.finditer(
r'<img\s[^>]*src=\"([^\"]+)\"',
markdown_text,
re.IGNORECASE,
):
raw_path = m.group(1).strip()
if not raw_path:
continue
if raw_path.startswith("http"):
url = raw_path
else:
clean = raw_path.lstrip("./").lstrip("/")
url = f"{base_url}/{clean}"
if url and url not in seen_urls:
seen_urls.add(url)
images.append({
"url": url,
"type": "image",
"nsfwLevel": 0,
"width": default_width,
"height": default_height,
"meta": {"prompt": "", "negativePrompt": ""},
"hasMeta": False,
"hasPositivePrompt": False,
})
# Also try single-quoted src attributes
for m in re.finditer(
r"<img\s[^>]*src='([^']+)'",
markdown_text,
re.IGNORECASE,
):
raw_path = m.group(1).strip()
if not raw_path:
continue
if raw_path.startswith("http"):
url = raw_path
else:
clean = raw_path.lstrip("./").lstrip("/")
url = f"{base_url}/{clean}"
if url and url not in seen_urls:
seen_urls.add(url)
images.append({
"url": url,
"type": "image",
"nsfwLevel": 0,
"width": default_width,
"height": default_height,
"meta": {"prompt": "", "negativePrompt": ""},
"hasMeta": False,
"hasPositivePrompt": False,
})
return images
def extract_repo_from_hf_url(hf_url: str) -> str:
"""Extract ``user/repo`` from a HuggingFace URL."""
m = _REPO_URL_PATTERN.match(hf_url)
@@ -62,43 +211,20 @@ def extract_gallery_images(
if not widget_match:
return images
# Split entries starting with `- text:`
entries = re.split(r"\n- text:", frontmatter[widget_match.end():])
# Split entries by YAML list marker `\n- `. Each entry is one widget list
# item with `output:` and optionally `text:` sub-keys.
entries_raw = frontmatter[widget_match.end():]
entries = re.split(r"\n- ", entries_raw) if "\n- " in entries_raw else [entries_raw]
for entry in entries:
if not entry.strip():
entry = entry.strip()
if not entry:
continue
entry = entry.strip()
# Extract text (prompt) — handles three YAML scalar styles:
# 1. "quoted inline"
# 2. >- folded block
# 3. plain (unquoted) multi-line
text = ""
qm = re.match(r'^"((?:[^"\\]|\\.)*)"', entry)
if qm:
text = qm.group(1)
else:
mm = re.match(r"^>(?:-\s*)?\n((?:.+(?:\n|$))+)", entry, re.MULTILINE)
if mm:
raw = mm.group(1)
else:
# Plain YAML scalar — take lines until a YAML key
raw = entry
if raw:
text_lines: list[str] = []
for line in raw.split("\n"):
if re.match(r"^\s*\w+:", line):
break
stripped = line.strip()
if stripped:
text_lines.append(stripped)
text = " ".join(text_lines)
# Extract output.url
# Extract output.url — handle both `output:` and `- output:` (dash prefix)
url = ""
url_match = re.search(
r"^\s*output:\s*\n\s+url:\s*(.+?)\s*$", entry, re.MULTILINE
r"^-?\s*output:\s*\n\s+url:\s*(.+?)\s*$", entry, re.MULTILINE
)
if url_match:
raw_path = url_match.group(1).strip().strip("'\"")
@@ -107,6 +233,16 @@ def extract_gallery_images(
elif raw_path.startswith("http"):
url = raw_path
# Extract text (prompt) — from YAML `text:` sub-key
text = ""
text_match = re.search(
r"^-?\s*text:\s*(.+?)\s*$", entry, re.MULTILINE
)
if text_match:
raw_text = text_match.group(1).strip().strip("'\"")
if raw_text and raw_text != "-":
text = raw_text
if url:
image: dict = {
"url": url,
@@ -338,20 +474,23 @@ def _strip_fenced_code_blocks(text: str) -> str:
def _strip_standalone_images(text: str) -> str:
"""Strip image embeds that occupy their own line.
"""Strip/compress image embeds for LLM-prompt injection.
Preserves the alt text from markdown images (``![alt](url)`` → ``alt``)
since it often describes what the model generates, which is useful signal
for tag/description extraction.
Markdown images (``![alt](url)``) are **kept intact** so the LLM can
extract a ``preview_url`` from them. Only the alt text is needed for
content signal; the URL is needed for image extraction.
HTML ``<img>`` tags on their own line are replaced by their alt text
(if any) or removed, since the LLM has difficulty extracting URLs from
raw HTML attributes.
"""
# Markdown: ``![alt](url)`` on its own line → keep alt text
# HTML: ``<img src="..." alt="..." ...>`` on its own line → keep alt text
text = re.sub(
r"^\s*!\[([^\]]*)\]\([^)]+\)\s*$",
r'^\s*<img\s[^>]*alt="([^"]*)"[^>]*/?>(?:</img>)?\s*$',
r"\1",
text,
flags=re.MULTILINE,
flags=re.MULTILINE | re.IGNORECASE,
)
# HTML: ``<img src="..." ...>`` on its own line → remove entirely
text = re.sub(
r'^\s*<img\s[^>]+/?>(?:</img>)?\s*$',
"",
@@ -479,6 +618,155 @@ def _strip_massive_lists(text: str) -> str:
return "\n".join(out)
def extract_relevant_section(
readme_content: str,
model_basename: str,
*,
context_lines: int = 30,
) -> str:
"""Find the section of a HuggingFace README relevant to a specific model file.
Many HF repos bundle multiple model files (LoRA collections) in a single
repo. This function trims the README to only the portion that references
the given *model_basename* (e.g. ``lora_zimage_myjs_turbo_beta01``),
dramatically reducing prompt length and focusing the LLM on the correct
metadata.
Matching strategies, tried in order:
1. **Download link** — a markdown or HTML link whose URL contains the
basename. Returns the surrounding block (previous heading → next
heading).
2. **Anchor ID** — an ``<a id="...">`` whose ``id`` matches a token from
the basename (split on ``_``/``-``).
3. **Section heading** — an ``<h1>``-``<h4>`` or markdown heading whose
text overlaps with tokens from the basename.
4. **Fallback** — the full README unchanged.
Args:
readme_content: Raw README.md content.
model_basename: Model file basename *without* extension (e.g.
``lora_zimage_myjs_turbo_beta01``).
context_lines: Number of lines of context before/after a matched
download link to include (default 30).
Returns:
Trimmed README text, or the original when no matching section is
found.
"""
if not readme_content or not model_basename:
return readme_content or ""
lines = readme_content.split("\n")
n = len(lines)
basename_lower = model_basename.lower()
# Tokens from the basename split on common separators
tokens = {t for t in re.split(r"[_\-.\s]+", basename_lower) if len(t) > 2}
# ------------------------------------------------------------------
# Strategy 1: Find a download link containing the basename
# ------------------------------------------------------------------
for idx, line in enumerate(lines):
# Match markdown links: [text](url) and HTML links: <a href="url">
# whose URL contains the basename.
if basename_lower in line.lower() and _looks_like_download_link(line):
return _extract_section(lines, idx, context_lines)
# ------------------------------------------------------------------
# Strategy 2: Find an anchor ID matching a token
# ------------------------------------------------------------------
for idx, line in enumerate(lines):
m = re.search(r'<a\s+id="([^"]+)"', line, re.IGNORECASE)
if m:
aid = m.group(1).lower()
if any(token in aid for token in tokens):
return _extract_section(lines, idx, context_lines)
# ------------------------------------------------------------------
# Strategy 3: Find an HTML or markdown heading with overlapping tokens
# ------------------------------------------------------------------
for idx, line in enumerate(lines):
# HTML heading: <h1...>, <h2...>, <h3...>, <h4...>
hm = re.search(r'<h[1-4][^>]*>(.+?)</h[1-4]>', line, re.IGNORECASE)
heading_text = ""
if hm:
heading_text = hm.group(1)
else:
# Markdown heading: ## text
mm = re.match(r"^#{1,4}\s+(.+?)\s*#*$", line)
if mm:
heading_text = mm.group(1)
if heading_text:
heading_lower = heading_text.lower()
# Check if any token appears in the heading
if any(token in heading_lower for token in tokens):
return _extract_section(lines, idx, context_lines)
# ------------------------------------------------------------------
# Fallback: return FULL readme
# ------------------------------------------------------------------
return readme_content
def _looks_like_download_link(line: str) -> bool:
"""Heuristic: does *line* look like it contains a model-file download link?"""
# Markdown: [text](url)
if re.search(r'\[([^\]]*)\]\(([^)]*\.safetensors[^)]*)\)', line):
return True
# Markdown: [text](url) containing "download" or "resolve"
if re.search(r'\[([^\]]*)\]\(([^)]*/(resolve|download)/[^)]*)\)', line):
return True
# HTML <a href="...safetensors">
if re.search(r'<a\s[^>]*href="[^"]*\.safetensors', line, re.IGNORECASE):
return True
return False
def _extract_section(
lines: list[str], match_idx: int, context_lines: int,
) -> str:
"""Return a window of *lines* around *match_idx* bounded by headings.
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.
"""
n = len(lines)
# Determine start — if match is a heading, start right there
if _is_heading(lines[match_idx]):
start = match_idx
else:
# Walk backward to find the nearest heading
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:
start = 0
break
if _is_heading(lines[i]):
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]):
end = i
break
return "\n".join(lines[start:end])
def _is_heading(line: str) -> bool:
"""Return True if *line* is a markdown or HTML heading (not a closing tag)."""
stripped = line.strip()
if re.match(r"^#{1,4}\s", stripped):
return True
if re.match(r"^<h[1-4](?:\s|>)", stripped, re.IGNORECASE):
return True
return False
def _compress_blank_lines(text: str) -> str:
"""Collapse runs of 3+ blank lines down to 2."""
return re.sub(r"\n{3,}", "\n\n", text)

96
py/services/llm_service.py
View File

@@ -551,7 +551,103 @@ class LLMService:
try:
return json.loads(content)
except (json.JSONDecodeError, TypeError) as parse_err:
# Last resort: attempt to salvage partial JSON (closing unclosed
# brackets/braces, truncating incomplete strings, etc.)
salvaged = _try_salvage_json(content)
if salvaged is not None:
logger.warning(
"LLM JSON salvaged from partial content (%d chars raw)",
len(content),
)
return salvaged
raise LLMResponseError(
f"LLM response could not be parsed as JSON after retry: {parse_err}\n"
f"Raw content: {content[:500]}"
) from parse_err
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
"""Attempt to repair and parse a truncated JSON string.
Handles common truncation patterns:
* Incomplete string value at the end (``"foo`` → ``"foo"``)
* Missing closing ``}`` or ``]`` (respecting nesting order)
* Trailing comma before closing bracket
* Extra text after the JSON object (e.g. markdown fences)
Returns the parsed dict on success, ``None`` if repair is impossible.
"""
if not raw:
return None
text = raw.strip()
# Strip markdown fences if the LLM wrapped the JSON
if text.startswith("```"):
end = text.find("\n")
text = text[end + 1:] if end != -1 else text[3:]
if text.endswith("```"):
text = text[:-3].rstrip()
# Find the first '{' and strip everything before it
start = text.find("{")
if start == -1:
return None
text = text[start:]
# Try to close an incomplete string at the end (e.g. ``"https://huggingf``)
# Pattern: ends mid-string (last quote is open)
if text.count('"') % 2 == 1:
text += '"'
# Ensure trailing commas before closing braces work
text = _strip_trailing_commas(text)
# Walk through the text character by character to find unclosed
# brackets and close them in the correct (LIFO) order.
# We ignore brackets inside quoted strings.
stack: list[str] = []
in_string = False
escape = False
for ch in text:
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if ch == '"':
in_string = not in_string
continue
if in_string:
continue
if ch in ("{", "["):
stack.append(ch)
elif ch == "}":
if stack and stack[-1] == "{":
stack.pop()
else:
return None # Unmatched closer — unrecoverable
elif ch == "]":
if stack and stack[-1] == "[":
stack.pop()
else:
return None
# Close remaining open brackets in reverse order
for opener in reversed(stack):
text += "}" if opener == "{" else "]"
try:
return json.loads(text)
except (json.JSONDecodeError, ValueError):
return None
def _strip_trailing_commas(text: str) -> str:
"""Remove commas that appear before a closing brace/bracket."""
import re as _re
text = _re.sub(r",\s*}", "}", text)
text = _re.sub(r",\s*]", "]", text)
return text