391 lines
14 KiB
Python
391 lines
14 KiB
Python
"""Per-sub-intent canonical-tag resolver.
|
|
|
|
Replaces the bge-only retrieval bottleneck where bare-form Danbooru
|
|
tags (`sitting`, `standing`, `barefoot`, `dojo`) were buried under
|
|
compound variants (`sitting_on_arm`, `holding_microphone_stand`) in
|
|
bge's top-K. The resolver:
|
|
|
|
1. Asks the configured LLM to propose 3-8 canonical Danbooru tags
|
|
for each sub-intent (in parallel across intents).
|
|
2. Validates each proposal against `danbooru_tags` (real-tag check
|
|
+ body-part variant swap, e.g. `feet_focus` -> `foot_focus`).
|
|
3. Surfaces validated tags as high-priority candidates with
|
|
`matched_via="resolved"`, which the patch user message routes
|
|
into the "Direct text matches" block.
|
|
|
|
Achieves ~86% resolution on the test corpus vs ~18% for bge alone.
|
|
The existing `tag_search` retrieval still runs — its results land in
|
|
the "Related candidates" block as before, so the resolver is purely
|
|
additive (cannot regress).
|
|
|
|
Skips out-of-scope sections (character / clear / style) — those go
|
|
through the bio system / template path / no-tag path. Cost is one
|
|
LLM call per in-scope sub-intent, parallelized; total wall-clock is
|
|
roughly the slowest single call (~0.6-0.8s on 8b)."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import re
|
|
from typing import Callable
|
|
|
|
from . import ai_api
|
|
|
|
logger = logging.getLogger("promptchain.canonical_resolver")
|
|
dbg = logging.getLogger("promptchain.ai.debug")
|
|
|
|
|
|
_RESOLVE_SKIP_SECTIONS = frozenset({"character", "clear", "style"})
|
|
|
|
|
|
_PROPOSE_SYSTEM = (
|
|
"You are a Danbooru tag expert. Given a natural-language phrase, "
|
|
"list 3-8 canonical Danbooru tags that could match. Output ONE "
|
|
"TAG PER LINE in lowercase underscore form (e.g. `foot_focus`). "
|
|
"No explanations, no markdown bullets, no punctuation, no header. "
|
|
"If you don't know any matching tags, output `NONE`.\n/no_think"
|
|
)
|
|
|
|
|
|
# Body-part word variants — the LLM occasionally emits the wrong number
|
|
# (`feet_focus` instead of `foot_focus`). One swap is allowed before the
|
|
# proposal is rejected.
|
|
_BODY_PART_PAIRS = [
|
|
("feet", "foot"), ("hands", "hand"), ("eyes", "eye"),
|
|
("fingers", "finger"), ("toes", "toe"), ("knees", "knee"),
|
|
("legs", "leg"), ("ears", "ear"), ("breasts", "breast"),
|
|
("thighs", "thigh"),
|
|
]
|
|
|
|
|
|
def _tag_exists(tag: str) -> bool:
|
|
"""True iff `tag` is in `danbooru_tags`."""
|
|
if not tag:
|
|
return False
|
|
try:
|
|
from .tag_builder import get_db
|
|
db = get_db()
|
|
row = db.execute(
|
|
"SELECT 1 FROM danbooru_tags WHERE tag = ? LIMIT 1",
|
|
(tag.lower().strip(),),
|
|
).fetchone()
|
|
return bool(row)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _try_body_part_swap(tag: str) -> str | None:
|
|
"""If `tag` isn't real but a single body-part swap (feet↔foot etc.)
|
|
yields a real tag, return that variant. Else None."""
|
|
parts = tag.split("_")
|
|
for i, p in enumerate(parts):
|
|
for plural, singular in _BODY_PART_PAIRS:
|
|
for current, swap_to in ((plural, singular), (singular, plural)):
|
|
if p != current:
|
|
continue
|
|
candidate = "_".join(parts[:i] + [swap_to] + parts[i + 1:])
|
|
if _tag_exists(candidate):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _validate(tag: str) -> str | None:
|
|
"""Return real canonical form (with body-part variant swap) or None."""
|
|
tag = (tag or "").strip().lower().replace(" ", "_")
|
|
if not tag:
|
|
return None
|
|
if _tag_exists(tag):
|
|
return tag
|
|
return _try_body_part_swap(tag)
|
|
|
|
|
|
def _parse_proposed_tags(raw: str) -> list[str]:
|
|
"""Parse the LLM's proposal output (one tag per line) into a list."""
|
|
if not raw or raw.strip().lower() == "none":
|
|
return []
|
|
out: list[str] = []
|
|
seen: set[str] = set()
|
|
for line in raw.splitlines():
|
|
line = re.sub(r"^[\s\-\*•·]*", "", line).strip()
|
|
line = re.sub(r"[`'\"]+|[\.,;:!?]$", "", line).strip()
|
|
if not line or line.lower() == "none":
|
|
continue
|
|
tok = line.split()[0].strip().lower().rstrip(".,;:!?")
|
|
if tok and tok not in seen:
|
|
seen.add(tok)
|
|
out.append(tok)
|
|
return out
|
|
|
|
|
|
async def _llm_propose(phrase: str, request_id: str) -> str:
|
|
"""Provider-agnostic LLM call returning the raw response text."""
|
|
config = ai_api._load_config()
|
|
provider = config.get("provider")
|
|
user_text = f'Phrase: "{phrase}"'
|
|
try:
|
|
return await ai_api._call_provider_complete(
|
|
request_id, provider, config, _PROPOSE_SYSTEM, user_text, [],
|
|
)
|
|
except Exception:
|
|
logger.warning("canonical_resolver: LLM propose call failed",
|
|
exc_info=True)
|
|
return ""
|
|
|
|
|
|
# Stopwords for the secondary fallback's ngram permutation. Prepositions
|
|
# and articles that connect content words but don't carry meaning.
|
|
# Mirrors the literal-anchor stopword set in ai_api but kept local so
|
|
# the resolver doesn't depend on private-named ai_api internals.
|
|
_NGRAM_STOPWORDS = frozenset({
|
|
"a", "an", "the", "her", "his", "their", "its", "is", "are",
|
|
"was", "were", "be", "to", "of", "and", "or", "on", "in", "at",
|
|
"for", "from", "with", "this", "that",
|
|
})
|
|
|
|
|
|
def _ngram_permute_validated(phrase: str) -> list[str]:
|
|
"""Deterministic word-permutation candidate generator. Splits phrase
|
|
into content words, generates 1-3-word permutations (ordered + body-
|
|
part variant swaps), validates each against danbooru_tags. Returns
|
|
the validated real-tag list, deduped, preserving discovery order.
|
|
|
|
Used by `_resolve_one` as a SECONDARY fallback when the primary LLM-
|
|
propose path returns empty. Catches `<focus|emphasis|frame>_<body
|
|
_part>` and similar reorderings: 'focus on feet' permutes through
|
|
`feet_focus` which body-part-swaps to the real `foot_focus`."""
|
|
if not phrase:
|
|
return []
|
|
words = [w for w in re.findall(r"[a-z]+", phrase.lower())
|
|
if w not in _NGRAM_STOPWORDS and len(w) >= 3]
|
|
if not words:
|
|
return []
|
|
candidates: list[str] = []
|
|
for w in words:
|
|
candidates.append(w)
|
|
for i in range(len(words)):
|
|
for j in range(len(words)):
|
|
if i == j:
|
|
continue
|
|
candidates.append(f"{words[i]}_{words[j]}")
|
|
if len(words) >= 3:
|
|
for i in range(len(words)):
|
|
for j in range(len(words)):
|
|
for k in range(len(words)):
|
|
if len({i, j, k}) < 3:
|
|
continue
|
|
candidates.append(f"{words[i]}_{words[j]}_{words[k]}")
|
|
out: list[str] = []
|
|
seen: set[str] = set()
|
|
for c in candidates:
|
|
if c in seen:
|
|
continue
|
|
seen.add(c)
|
|
v = _validate(c)
|
|
if v and v not in seen:
|
|
seen.add(v)
|
|
out.append(v)
|
|
return out
|
|
|
|
|
|
_RERANK_SYSTEM = (
|
|
"You pick the single canonical Danbooru tag that best matches a "
|
|
"user's phrase, from a curated list of REAL tags. Output ONLY the "
|
|
"tag in lowercase underscore form (e.g. `foot_focus`) — no "
|
|
"explanation, no markdown, no punctuation. If none of the tags "
|
|
"capture the phrase's meaning, output `NONE`.\n/no_think"
|
|
)
|
|
|
|
|
|
async def _llm_rerank(phrase: str, candidates: list[str],
|
|
request_id: str) -> str | None:
|
|
"""Ask the LLM to pick the best-matching tag from a curated real-tag
|
|
list. Returns the picked tag (lowercased, underscored) or None when
|
|
the LLM declined or the response didn't validate."""
|
|
if not candidates:
|
|
return None
|
|
candidate_list = ", ".join(candidates)
|
|
user_text = (
|
|
f'Phrase: "{phrase}"\n\nReal Danbooru tags found by word '
|
|
f'permutation: {candidate_list}\n\nWhich one best matches the '
|
|
f"phrase's meaning? Output the tag only."
|
|
)
|
|
try:
|
|
raw = await ai_api._call_provider_complete(
|
|
request_id, ai_api._load_config().get("provider"),
|
|
ai_api._load_config(),
|
|
_RERANK_SYSTEM, user_text, [],
|
|
)
|
|
except Exception:
|
|
logger.warning("canonical_resolver: LLM rerank call failed",
|
|
exc_info=True)
|
|
return None
|
|
raw = (raw or "").strip()
|
|
if not raw or raw.lower() == "none":
|
|
return None
|
|
pick = re.sub(r"[`'\"]+|[\.,;:!?]$", "", raw).strip()
|
|
pick = pick.split()[0].lower().rstrip(".,;:!?") if pick.split() else ""
|
|
if not pick:
|
|
return None
|
|
# Sanity: only return picks that are actually in the candidate list
|
|
# (prevents the LLM from reverting to a hallucination on rerank).
|
|
candidates_lc = {c.lower() for c in candidates}
|
|
return pick if pick in candidates_lc else None
|
|
|
|
|
|
def _build_resolved_candidates(validated: list[str], phrase: str,
|
|
via: str = "resolved") -> list[dict]:
|
|
"""Build the candidate-dict list for a list of validated tag forms.
|
|
Pulls wiki rows in one DB query so each candidate carries
|
|
body_summary for the patch user message."""
|
|
if not validated:
|
|
return []
|
|
rows: dict[str, dict] = {}
|
|
try:
|
|
from .tag_builder import get_db
|
|
db = get_db()
|
|
placeholders = ",".join("?" for _ in validated)
|
|
for r in db.execute(
|
|
f"SELECT t.tag, t.ranking, w.body_summary, w.body_full "
|
|
f"FROM danbooru_tags t "
|
|
f"LEFT JOIN danbooru_tag_wikis w ON w.tag = t.tag "
|
|
f"WHERE t.tag IN ({placeholders})",
|
|
validated,
|
|
).fetchall():
|
|
tag = (r["tag"] or "").lower()
|
|
rows[tag] = {
|
|
"tag": r["tag"],
|
|
"ranking": int(r["ranking"] or 0),
|
|
"body_summary": r["body_summary"] or "",
|
|
"body_full": r["body_full"] or "",
|
|
}
|
|
except Exception:
|
|
logger.warning("canonical_resolver: wiki lookup failed", exc_info=True)
|
|
out: list[dict] = []
|
|
for v in validated:
|
|
meta = rows.get(v, {
|
|
"tag": v, "ranking": 0,
|
|
"body_summary": "", "body_full": "",
|
|
})
|
|
out.append({
|
|
"tag": meta["tag"],
|
|
"ranking": meta["ranking"],
|
|
"score": 1.0, # sentinel — resolved tags always survive cap
|
|
"body_summary": meta["body_summary"],
|
|
"body_full": meta["body_full"],
|
|
"matched_intent": phrase,
|
|
"matched_via": via,
|
|
})
|
|
return out
|
|
|
|
|
|
async def _resolve_one(phrase: str, request_id: str) -> list[dict]:
|
|
"""Resolve one sub-intent phrase to candidate dicts.
|
|
|
|
Two-stage:
|
|
1. Primary — LLM proposes 3-8 tags from world knowledge,
|
|
validate against danbooru_tags + body-part variant swap.
|
|
2. Secondary (fires only when primary returns nothing) — generate
|
|
word permutations from the phrase, validate each, then ask the
|
|
LLM to rerank from that curated real-tag list. Catches
|
|
`<focus|frame|emphasis>_<body_part>` cases where the LLM's
|
|
proposal stage flaked but a permutation of the user's words
|
|
yields the canonical via body-part variant swap.
|
|
|
|
Secondary cost: one extra LLM call only when primary fails (~15% of
|
|
sub-intents on 8b). Parallelization across sub-intents bounds the
|
|
per-prompt latency hit to ~0.6s on prompts where any sub-intent's
|
|
primary flaked."""
|
|
# Primary path
|
|
raw = await _llm_propose(phrase, request_id)
|
|
proposed = _parse_proposed_tags(raw)
|
|
validated: list[str] = []
|
|
seen: set[str] = set()
|
|
for p in proposed:
|
|
v = _validate(p)
|
|
if v and v not in seen:
|
|
seen.add(v)
|
|
validated.append(v)
|
|
if validated:
|
|
return _build_resolved_candidates(validated, phrase, via="resolved")
|
|
|
|
# Secondary path — primary returned no real tags. Try ngram
|
|
# permutation against danbooru_tags (with body-part variant swap),
|
|
# then LLM rerank to pick the best match.
|
|
permuted = _ngram_permute_validated(phrase)
|
|
if not permuted:
|
|
return []
|
|
pick = await _llm_rerank(phrase, permuted, f"{request_id}-rerank")
|
|
if not pick:
|
|
return []
|
|
dbg.info(
|
|
"canonical_resolver[%s] secondary fallback fired: phrase=%r "
|
|
"permuted=%s pick=%s",
|
|
request_id, phrase, permuted[:8], pick,
|
|
)
|
|
return _build_resolved_candidates([pick], phrase, via="resolved-rerank")
|
|
|
|
|
|
async def resolve_intents_parallel(
|
|
sub_intents: list[dict],
|
|
request_id: str,
|
|
on_status: Callable[[str], None] | None = None,
|
|
) -> list[dict]:
|
|
"""Resolve every in-scope sub-intent in parallel. Returns a flat,
|
|
deduped list of candidate dicts suitable for merging into the
|
|
patch flow's existing `tag_candidates` list. Empty list on no
|
|
in-scope intents or total resolution failure."""
|
|
if not sub_intents:
|
|
return []
|
|
|
|
in_scope: list[dict] = []
|
|
for si in sub_intents:
|
|
section = (si.get("section") or "").lower()
|
|
text = (si.get("text") or "").strip()
|
|
if not text or section in _RESOLVE_SKIP_SECTIONS:
|
|
continue
|
|
in_scope.append(si)
|
|
if not in_scope:
|
|
return []
|
|
|
|
if on_status:
|
|
on_status("Resolving canonical tags")
|
|
|
|
tasks = [
|
|
_resolve_one(si["text"], f"{request_id}-resolve-{i}")
|
|
for i, si in enumerate(in_scope)
|
|
]
|
|
try:
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
except Exception:
|
|
logger.exception("canonical_resolver: parallel resolution failed")
|
|
return []
|
|
|
|
flat: list[dict] = []
|
|
seen: set[str] = set()
|
|
for si, result in zip(in_scope, results):
|
|
if isinstance(result, Exception):
|
|
logger.warning(
|
|
"canonical_resolver: intent %r failed: %s",
|
|
si.get("text"), result,
|
|
)
|
|
continue
|
|
for cand in result:
|
|
tag = (cand.get("tag") or "").lower()
|
|
if tag in seen:
|
|
continue
|
|
seen.add(tag)
|
|
flat.append(cand)
|
|
|
|
if flat:
|
|
dbg.info(
|
|
"canonical_resolver[%s] %d intents -> %d resolved tags: %s",
|
|
request_id, len(in_scope), len(flat),
|
|
", ".join(c["tag"] for c in flat[:12]),
|
|
)
|
|
else:
|
|
dbg.info(
|
|
"canonical_resolver[%s] %d intents -> 0 resolved tags",
|
|
request_id, len(in_scope),
|
|
)
|
|
return flat
|