| 1 | """Runtime-switchable Prompt Engineering (PE) sets. |
| 2 | |
| 3 | A *PE set* is a directory under the configured ``pe/`` root: |
| 4 | |
| 5 | pe/<name>/ |
| 6 | manifest.yaml # {name, label, description} |
| 7 | strings.yaml # extracted prompt strings, dotted keys after flattening |
| 8 | templates/ # optional Jinja template overrides (e.g. agent/identity.md) |
| 9 | skills/<skill>/SKILL.md # optional skill markdown overrides |
| 10 | bootstrap/SOUL.md # optional workspace-bootstrap overrides (SOUL/TOOLS/...) |
| 11 | |
| 12 | The active set overlays ``default`` for every resource: strings fall back to |
| 13 | ``default`` per key; template/skill/bootstrap lookups search ``active`` then |
| 14 | ``default`` then the packaged locations. |
| 15 | |
| 16 | ``PEManager`` is a process-wide singleton so the WebUI can hot-switch the active |
| 17 | set without a restart. Consumers that cache derived artifacts (e.g. the tool |
| 18 | schema cache) subscribe via :meth:`on_change`. |
| 19 | """ |
| 20 | |
| 21 | from __future__ import annotations |
| 22 | |
| 23 | import logging |
| 24 | import threading |
| 25 | from pathlib import Path |
| 26 | from typing import Any, Callable |
| 27 | |
| 28 | import yaml |
| 29 | |
| 30 | logger = logging.getLogger(__name__) |
| 31 | |
| 32 | DEFAULT_SET_NAME = "default" |
| 33 | |
| 34 | |
| 35 | def default_pe_root() -> Path: |
| 36 | """Return source-checkout PE resources, or the resources bundled in a wheel.""" |
| 37 | source_root = Path(__file__).resolve().parents[2] / "pe" |
| 38 | if source_root.is_dir(): |
| 39 | return source_root |
| 40 | return Path(__file__).resolve().parents[1] / "_bundled" / "pe" |
| 41 | |
| 42 | |
| 43 | def _flatten(data: Any, prefix: str = "") -> dict[str, str]: |
| 44 | """Flatten a nested mapping into dotted keys with string values.""" |
| 45 | out: dict[str, str] = {} |
| 46 | if isinstance(data, dict): |
| 47 | for key, value in data.items(): |
| 48 | child = f"{prefix}.{key}" if prefix else str(key) |
| 49 | if isinstance(value, dict): |
| 50 | out.update(_flatten(value, child)) |
| 51 | elif value is not None: |
| 52 | out[child] = value if isinstance(value, str) else str(value) |
| 53 | return out |
| 54 | |
| 55 | |
| 56 | class PEManager: |
| 57 | """Process-wide holder of the active PE set and resource-resolution helpers.""" |
| 58 | |
| 59 | _instance: "PEManager | None" = None |
| 60 | _instance_lock = threading.Lock() |
| 61 | |
| 62 | def __init__(self) -> None: |
| 63 | self._root: Path | None = None |
| 64 | self._active: str = DEFAULT_SET_NAME |
| 65 | self._enabled: bool = True |
| 66 | self._listeners: list[Callable[[str], None]] = [] |
| 67 | self._strings_cache: dict[str, dict[str, str]] = {} |
| 68 | # Per-session active-set overrides (in-memory; lost on restart). |
| 69 | self._session_sets: dict[str, str] = {} |
| 70 | |
| 71 | @classmethod |
| 72 | def instance(cls) -> "PEManager": |
| 73 | if cls._instance is None: |
| 74 | with cls._instance_lock: |
| 75 | if cls._instance is None: |
| 76 | inst = cls() |
| 77 | root = default_pe_root() |
| 78 | if root.is_dir(): |
| 79 | inst.configure(root) |
| 80 | cls._instance = inst |
| 81 | return cls._instance |
| 82 | |
| 83 | # ------------------------------------------------------------------ config |
| 84 | |
| 85 | def configure( |
| 86 | self, |
| 87 | root: str | Path | None, |
| 88 | active: str = DEFAULT_SET_NAME, |
| 89 | enabled: bool = True, |
| 90 | ) -> None: |
| 91 | """(Re)initialize from config. Called once at CLI startup.""" |
| 92 | self._root = Path(root).expanduser().resolve() if root else None |
| 93 | self._enabled = enabled |
| 94 | self._strings_cache.clear() |
| 95 | names = {entry["name"] for entry in self.list_sets()} |
| 96 | if active in names: |
| 97 | self._active = active |
| 98 | elif DEFAULT_SET_NAME in names: |
| 99 | self._active = DEFAULT_SET_NAME |
| 100 | else: |
| 101 | self._active = active |
| 102 | logger.info( |
| 103 | "PEManager configured: root=%s active=%s enabled=%s available=%s", |
| 104 | self._root, |
| 105 | self._active, |
| 106 | self._enabled, |
| 107 | sorted(names), |
| 108 | ) |
| 109 | |
| 110 | @property |
| 111 | def root(self) -> Path | None: |
| 112 | return self._root |
| 113 | |
| 114 | @property |
| 115 | def active(self) -> str: |
| 116 | return self._active |
| 117 | |
| 118 | @property |
| 119 | def enabled(self) -> bool: |
| 120 | return self._enabled |
| 121 | |
| 122 | # -------------------------------------------------------------------- sets |
| 123 | |
| 124 | def list_sets(self) -> list[dict[str, str]]: |
| 125 | """Return metadata for every PE set directory under the root.""" |
| 126 | out: list[dict[str, str]] = [] |
| 127 | if not self._root or not self._root.is_dir(): |
| 128 | return out |
| 129 | for child in sorted(self._root.iterdir()): |
| 130 | if not child.is_dir(): |
| 131 | continue |
| 132 | meta = {"name": child.name, "label": child.name, "description": ""} |
| 133 | manifest = child / "manifest.yaml" |
| 134 | if manifest.is_file(): |
| 135 | try: |
| 136 | data = yaml.safe_load(manifest.read_text(encoding="utf-8")) or {} |
| 137 | if data.get("name"): |
| 138 | meta["name"] = str(data["name"]) |
| 139 | meta["label"] = str(data.get("label") or meta["name"]) |
| 140 | meta["description"] = str(data.get("description") or "") |
| 141 | except Exception as exc: |
| 142 | logger.warning("Failed to parse PE manifest %s: %s", manifest, exc) |
| 143 | out.append(meta) |
| 144 | return out |
| 145 | |
| 146 | def set_active(self, name: str) -> bool: |
| 147 | """Switch the active set and notify subscribers. Returns success.""" |
| 148 | if not self._enabled: |
| 149 | logger.warning("PE switching disabled; ignoring set_active(%r)", name) |
| 150 | return False |
| 151 | names = {entry["name"] for entry in self.list_sets()} |
| 152 | if name not in names: |
| 153 | logger.warning("PE set %r not found; ignoring", name) |
| 154 | return False |
| 155 | if name == self._active: |
| 156 | return True |
| 157 | self._active = name |
| 158 | logger.info("PE active set switched to %r", name) |
| 159 | self._notify() |
| 160 | return True |
| 161 | |
| 162 | def active_for_session(self, session_key: str | None) -> str: |
| 163 | """Resolve the active set for a session: per-session override, else global. |
| 164 | |
| 165 | A stale override (set no longer present on disk) falls back to global. |
| 166 | """ |
| 167 | if session_key: |
| 168 | override = self._session_sets.get(session_key) |
| 169 | if override and override in {e["name"] for e in self.list_sets()}: |
| 170 | return override |
| 171 | return self._active |
| 172 | |
| 173 | def set_active_for_session(self, session_key: str, name: str) -> bool: |
| 174 | """Bind ``name`` as the active set for ``session_key`` only. Returns success.""" |
| 175 | if not self._enabled: |
| 176 | logger.warning("PE switching disabled; ignoring set_active_for_session(%r)", name) |
| 177 | return False |
| 178 | if not session_key: |
| 179 | return False |
| 180 | if name not in {e["name"] for e in self.list_sets()}: |
| 181 | logger.warning("PE set %r not found; ignoring session override", name) |
| 182 | return False |
| 183 | self._session_sets[session_key] = name |
| 184 | logger.info("PE set %r bound to session %s", name, session_key) |
| 185 | return True |
| 186 | |
| 187 | def on_change(self, callback: Callable[[str], None]) -> None: |
| 188 | """Register a callback fired (with the new active name) on every switch.""" |
| 189 | self._listeners.append(callback) |
| 190 | |
| 191 | def _notify(self) -> None: |
| 192 | for callback in list(self._listeners): |
| 193 | try: |
| 194 | callback(self._active) |
| 195 | except Exception: |
| 196 | logger.exception("PE on_change listener failed") |
| 197 | |
| 198 | # ------------------------------------------------------ resource resolution |
| 199 | |
| 200 | def _resource_dirs(self, subdir: str, name: str | None = None) -> list[str]: |
| 201 | """Existing ``<set>/<subdir>`` dirs for active then default (deduped). |
| 202 | |
| 203 | ``name`` overrides the active set (used for per-session resolution). |
| 204 | """ |
| 205 | if not self._root: |
| 206 | return [] |
| 207 | dirs: list[str] = [] |
| 208 | seen: set[str] = set() |
| 209 | primary = name or self._active |
| 210 | candidates = [primary, DEFAULT_SET_NAME] if self._enabled else [primary] |
| 211 | for set_name in candidates: |
| 212 | path = self._root / set_name / subdir |
| 213 | key = str(path) |
| 214 | if key not in seen and path.is_dir(): |
| 215 | dirs.append(key) |
| 216 | seen.add(key) |
| 217 | return dirs |
| 218 | |
| 219 | def templates_dir(self) -> list[str]: |
| 220 | return self._resource_dirs("templates") |
| 221 | |
| 222 | def skills_dir(self) -> list[str]: |
| 223 | return self._resource_dirs("skills") |
| 224 | |
| 225 | def bootstrap_dir(self) -> list[str]: |
| 226 | return self._resource_dirs("bootstrap") |
| 227 | |
| 228 | def references_dir(self, name: str | None = None) -> list[str]: |
| 229 | return self._resource_dirs("references", name) |
| 230 | |
| 231 | def resolve_reference(self, topic: str, name: str | None = None) -> Path | None: |
| 232 | """Resolve a guidance topic to a ``<set>/references/<topic>.md`` file. |
| 233 | |
| 234 | Active set overlays default: the first existing match wins. ``name`` |
| 235 | overrides the active set (used for per-session resolution). |
| 236 | """ |
| 237 | stem = topic.strip().removesuffix(".md") |
| 238 | if not stem: |
| 239 | return None |
| 240 | for base in self.references_dir(name): |
| 241 | for candidate in (Path(base) / f"{stem}.md", Path(base) / stem): |
| 242 | if candidate.is_file(): |
| 243 | return candidate |
| 244 | return None |
| 245 | |
| 246 | def list_references(self, name: str | None = None) -> list[str]: |
| 247 | """Available guidance topics (``.md`` basenames), active overlaying default.""" |
| 248 | topics: dict[str, None] = {} |
| 249 | for base in self.references_dir(name): |
| 250 | for path in sorted(Path(base).glob("*.md")): |
| 251 | topics.setdefault(path.stem, None) |
| 252 | return list(topics) |
| 253 | |
| 254 | # -------------------------------------------------------------- strings |
| 255 | |
| 256 | def _load_strings(self, name: str) -> dict[str, str]: |
| 257 | if name in self._strings_cache: |
| 258 | return self._strings_cache[name] |
| 259 | flat: dict[str, str] = {} |
| 260 | if self._root: |
| 261 | path = self._root / name / "strings.yaml" |
| 262 | if path.is_file(): |
| 263 | try: |
| 264 | data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} |
| 265 | flat = _flatten(data) |
| 266 | except Exception as exc: |
| 267 | logger.warning("Failed to parse PE strings %s: %s", path, exc) |
| 268 | self._strings_cache[name] = flat |
| 269 | return flat |
| 270 | |
| 271 | def get_string(self, key: str) -> str | None: |
| 272 | """Resolve a flattened string key: active overlays default.""" |
| 273 | if not self._root: |
| 274 | return None |
| 275 | if self._enabled: |
| 276 | active = self._load_strings(self._active) |
| 277 | if key in active: |
| 278 | return active[key] |
| 279 | base = self._load_strings(DEFAULT_SET_NAME) |
| 280 | return base.get(key) |
| 281 | |
| 282 | def reload(self) -> None: |
| 283 | """Drop cached strings (e.g. after editing YAML on disk).""" |
| 284 | self._strings_cache.clear() |
| 285 |