digital-garbage
d80be3c327
A variable referencing the same other variable more than once is not a cycle, so memoization alone still leaves the OUTPUT exponential in size (each level's re.sub concatenates the memoized child result multiple times). Adds a shared output-size budget, checked before using either a cached or freshly-computed result. Also ignore the runtime shortcuts.json data file, and fix the fan-out test's assertion, which was checking against a vacuous condition.
283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""The shortcut library: trigger -> replacement text, expanded into a prompt at
|
||
generation time. Global, not per-project -- a shortcut ("cyberpunk" -> a long
|
||
style phrase) is reused across every project, not redefined in each one.
|
||
|
||
Storage is one JSON file holding the whole list, not one-per-item: the picker
|
||
browses all of it and export/import round-trips all of it, so there is no
|
||
case where only part of the library is ever read or written.
|
||
|
||
The expansion algorithm (trigger regex, longest-first, seeded multi-choice) and
|
||
`$variable` resolution (recursive, cycle-safe, undefined -> literal) are a
|
||
close port of v4's `templates.py` (`apply_prompt_shortcuts`/`resolve_variables`)
|
||
-- proven logic, not reinvented. Not ported: the shortcut revolver's
|
||
no-repeat cycling (a stateful commit-vs-preview split) and managed category
|
||
lists (empty categories kept around for a picker to offer) -- both are real
|
||
v4 features, cut here to keep this a first pass rather than a rewrite of
|
||
the whole prompt-craft surface at once.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import random
|
||
import re
|
||
import threading
|
||
from dataclasses import asdict, dataclass, field
|
||
|
||
from . import config
|
||
|
||
_LOCK = threading.Lock()
|
||
|
||
MAX_NAME = 120
|
||
MAX_ITEM = 4096
|
||
|
||
|
||
def _clean_list(raw, *, keep_empty=False) -> list[str]:
|
||
"""One-per-line strings: trimmed, internal whitespace collapsed, exact
|
||
duplicates dropped. Triggers and replacements are both lists of whole
|
||
phrases -- comma-splitting a replacement tears one prose phrase into
|
||
bogus variants, so callers never do that upstream of here.
|
||
|
||
`keep_empty` is for replacements only: an empty string there is a
|
||
deliberate "remove this phrase" entry (`expand()`'s `_cleanup_removed_
|
||
phrases`), not a blank line to discard -- a trigger has no such meaning
|
||
and always drops empties."""
|
||
if isinstance(raw, str):
|
||
raw = raw.split("\n")
|
||
if not isinstance(raw, list):
|
||
return []
|
||
seen, out = set(), []
|
||
for item in raw:
|
||
text = re.sub(r"\s+", " ", str(item if item is not None else "")).strip()[:MAX_ITEM]
|
||
if (text or keep_empty) and text not in seen:
|
||
seen.add(text)
|
||
out.append(text)
|
||
return out
|
||
|
||
|
||
@dataclass
|
||
class Shortcut:
|
||
name: str = ""
|
||
triggers: list[str] = field(default_factory=list)
|
||
replacements: list[str] = field(default_factory=list)
|
||
enabled: bool = True
|
||
#: Free-text grouping for the picker UI. Plain strings, not a managed
|
||
#: list -- a category exists the moment a shortcut is saved under it, and
|
||
#: goes away the moment nothing uses it any more, no separate CRUD needed.
|
||
category: str = ""
|
||
sub_category: str = ""
|
||
|
||
@staticmethod
|
||
def from_dict(d) -> "Shortcut":
|
||
d = d if isinstance(d, dict) else {}
|
||
triggers = _clean_list(d.get("triggers"))
|
||
name = str(d.get("name") or "").strip()[:MAX_NAME] or (triggers[0] if triggers else "")
|
||
return Shortcut(
|
||
name=name,
|
||
triggers=triggers,
|
||
replacements=_clean_list(d.get("replacements"), keep_empty=True),
|
||
enabled=bool(d.get("enabled", True)),
|
||
category=str(d.get("category") or "").strip()[:MAX_NAME],
|
||
sub_category=str(d.get("sub_category") or "").strip()[:MAX_NAME],
|
||
)
|
||
|
||
def to_dict(self) -> dict:
|
||
return asdict(self)
|
||
|
||
|
||
def _path():
|
||
config.ROOT.mkdir(parents=True, exist_ok=True)
|
||
return config.SHORTCUTS_FILE
|
||
|
||
|
||
def listing() -> list[Shortcut]:
|
||
p = config.SHORTCUTS_FILE
|
||
if not p.exists():
|
||
return []
|
||
try:
|
||
data = json.loads(p.read_text(encoding="utf-8"))
|
||
except (json.JSONDecodeError, OSError, UnicodeDecodeError):
|
||
return [] # an unreadable library is not a reason to have no list
|
||
items = data if isinstance(data, list) else []
|
||
return [Shortcut.from_dict(it) for it in items]
|
||
|
||
|
||
def _save_all(items: list[Shortcut]) -> None:
|
||
path = _path()
|
||
tmp = path.with_suffix(".tmp")
|
||
tmp.write_text(json.dumps([it.to_dict() for it in items], indent=2), encoding="utf-8")
|
||
tmp.replace(path) # a concurrent read never sees a half-written library
|
||
|
||
|
||
def save(payload: dict, original_name: str | None = None) -> list[Shortcut]:
|
||
"""Insert, or update in place when `original_name` names an existing
|
||
entry -- a rename keeps its position; without this a save-under-a-new-
|
||
name would move the entry to the end, reordering the picker under the
|
||
user's feet for no reason they asked for. Raises ValueError for a
|
||
shortcut with no triggers: nothing can ever match it, so saving one
|
||
would be silent data loss dressed up as success."""
|
||
item = Shortcut.from_dict(payload)
|
||
if not item.triggers:
|
||
raise ValueError("a shortcut needs at least one trigger")
|
||
with _LOCK:
|
||
items = listing()
|
||
idx = next((i for i, it in enumerate(items)
|
||
if it.name == (original_name or item.name)), None)
|
||
if idx is None:
|
||
items.append(item)
|
||
else:
|
||
items[idx] = item
|
||
_save_all(items)
|
||
return items
|
||
|
||
|
||
def delete(name: str) -> list[Shortcut]:
|
||
with _LOCK:
|
||
items = [it for it in listing() if it.name != name]
|
||
_save_all(items)
|
||
return items
|
||
|
||
|
||
def clear() -> list[Shortcut]:
|
||
with _LOCK:
|
||
_save_all([])
|
||
return []
|
||
|
||
|
||
# --- shortcut expansion ------------------------------------------------------
|
||
|
||
def _trigger_pattern(trigger: str) -> str:
|
||
words = [re.escape(w) for w in re.split(r"\s+", trigger.strip()) if w]
|
||
if not words:
|
||
return ""
|
||
body = r"\s+".join(words)
|
||
return rf"(?<![\w'’-])({body})(?![\w'’-])"
|
||
|
||
|
||
def _cleanup_removed_phrases(text: str) -> str:
|
||
"""Fix punctuation/spacing left behind when a replacement is "" (a
|
||
deliberate 'remove this phrase' entry)."""
|
||
text = re.sub(r"[ \t]+([,;])", r"\1", text)
|
||
text = re.sub(r"([,;])\s*([,;])+", r"\1", text)
|
||
text = re.sub(r"^[\s,;]+", "", text)
|
||
text = re.sub(r"[\s,;]+$", "", text)
|
||
return re.sub(r"[ \t]{2,}", " ", text)
|
||
|
||
|
||
def expand(text: str, shortcuts: list[Shortcut] | None = None, seed: int = 0) -> str:
|
||
"""Every enabled shortcut's trigger, replaced with one of its
|
||
replacements (random when there is more than one). Deterministic: the
|
||
pick is seeded from `seed`, or from the text itself when seed is 0 --
|
||
so calling this twice on the same text and seed always agrees, without a
|
||
database round-trip to remember what was drawn last time.
|
||
|
||
Longest trigger wins first ("golden hour" is tried before "golden"), so
|
||
a short trigger can never shadow a longer one that contains it.
|
||
"""
|
||
original = str(text or "")
|
||
if not original:
|
||
return original
|
||
items = listing() if shortcuts is None else shortcuts
|
||
candidates = []
|
||
for sc in items:
|
||
if not sc.enabled or not sc.replacements:
|
||
continue
|
||
for trigger in sc.triggers:
|
||
pattern = _trigger_pattern(trigger)
|
||
if pattern:
|
||
candidates.append((trigger, pattern, sc.replacements))
|
||
if not candidates:
|
||
return original
|
||
|
||
candidates.sort(key=lambda c: len(c[0]), reverse=True)
|
||
combined = "|".join(f"(?P<t{i}>{p})" for i, (_, p, _) in enumerate(candidates))
|
||
rng_seed = int(seed or 0) or int(hashlib.md5(original.encode("utf-8")).hexdigest()[:12], 16)
|
||
rng = random.Random(rng_seed)
|
||
removed = False
|
||
|
||
def replace(m):
|
||
nonlocal removed
|
||
for i, (_, _, replacements) in enumerate(candidates):
|
||
if m.group(f"t{i}") is None:
|
||
continue
|
||
choice = rng.choice(replacements)
|
||
if not choice:
|
||
removed = True
|
||
return choice
|
||
return m.group(0) # unreachable: `combined` only matches a known group
|
||
|
||
expanded = re.sub(combined, replace, original, flags=re.IGNORECASE | re.UNICODE)
|
||
return _cleanup_removed_phrases(expanded) if removed else expanded
|
||
|
||
|
||
# --- $variables --------------------------------------------------------------
|
||
|
||
_VARIABLE_TOKEN = re.compile(r"\$([A-Za-z_][A-Za-z0-9_]*)")
|
||
_VARIABLE_MAX_DEPTH = 64
|
||
#: A ceiling on total resolved length, in characters -- far past anything a
|
||
#: real prompt needs, and not a working limit for one. It exists because a
|
||
#: variable referencing the same other variable more than once is not a
|
||
#: cycle (see below) and the OUTPUT for that shape is legitimately
|
||
#: exponential even once the per-(name,stack) evaluations are memoized: each
|
||
#: level's re.sub still has to concatenate the (now O(1)-computed) child
|
||
#: result multiple times, so the memo bounds compute but not output size.
|
||
_VARIABLE_MAX_OUTPUT = 100_000
|
||
|
||
|
||
def resolve_variables(text: str, variables) -> str:
|
||
"""Substitute `$name` tokens from a project's variables ([{"name",
|
||
"value"}, ...]). Recursive -- a variable's value may itself reference
|
||
another -- cycle-safe (a name reappearing in its own expansion chain is
|
||
left literal rather than recursed into), and undefined names are left as
|
||
literal `$name`: an unset variable is a typo to notice and fix, not a
|
||
blank the prompt swallows silently.
|
||
|
||
Memoized per (name, stack): a value referencing the SAME variable more
|
||
than once (`v0 = "$v1 $v1"`) is not a cycle -- `name in stack` never
|
||
trips -- so without this a chain of only two dozen such variables costs
|
||
2**24 evaluations of the last one and tens of megabytes of output before
|
||
_VARIABLE_MAX_DEPTH ever gets the chance to matter (that cap bounds
|
||
depth, not the branching this shape produces). The cache key includes
|
||
`stack`, not just `name`, because the SAME variable can legitimately
|
||
resolve differently depending on which cycle it is being asked from
|
||
(A="$B", B="$A": resolving from A ends in literal "$A", from B in
|
||
literal "$B") -- collapsing by name alone would answer one of those with
|
||
the other's result. Past memoization, `_VARIABLE_MAX_OUTPUT` bounds the
|
||
total characters produced -- once hit, remaining tokens are left literal
|
||
rather than expanded.
|
||
"""
|
||
var_map: dict[str, str] = {}
|
||
for v in (variables or []):
|
||
name = str((v or {}).get("name") or "").lstrip("$").strip()
|
||
if name:
|
||
var_map[name] = str((v or {}).get("value") or "")
|
||
if not var_map:
|
||
return str(text or "")
|
||
|
||
memo: dict[tuple[str, frozenset], str] = {}
|
||
budget = [_VARIABLE_MAX_OUTPUT]
|
||
|
||
def _expand(s, stack, depth):
|
||
if depth > _VARIABLE_MAX_DEPTH:
|
||
return s
|
||
|
||
def _repl(m):
|
||
if budget[0] <= 0:
|
||
return m.group(0)
|
||
name = m.group(1)
|
||
if name not in var_map or name in stack:
|
||
return m.group(0)
|
||
key = (name, stack)
|
||
if key in memo:
|
||
result = memo[key]
|
||
else:
|
||
result = _expand(var_map[name], stack | {name}, depth + 1)
|
||
memo[key] = result
|
||
budget[0] -= len(result)
|
||
return result
|
||
|
||
return _VARIABLE_TOKEN.sub(_repl, s)
|
||
|
||
return _expand(str(text or ""), frozenset(), 0)
|