wuenlp_tools.pipeline_catalog

  1"""Auto-discovered catalog of wuenlp-tools pipeline steps for declarative configuration."""
  2
  3from __future__ import annotations
  4
  5import importlib
  6import inspect
  7import json
  8import pkgutil
  9import re
 10from dataclasses import dataclass
 11from typing import Any
 12
 13from wuenlp_tools.pipeline import PipelineStep, PipelineCapability
 14
 15_STEP_FACTORY_NAME = re.compile(r".*(PipelineStep|Step)$")
 16
 17_SKIP_MODULE_PREFIXES = (
 18    "wuenlp_tools.pipeline",
 19    "wuenlp_tools.pipeline_catalog",
 20    "wuenlp_tools.run_pipeline",
 21    "wuenlp_tools.process_file",
 22    "wuenlp_tools.vendored",
 23    "wuenlp_tools.eval",
 24)
 25
 26_DISCOVERY_ROOTS = (
 27    "wuenlp_tools.models",
 28    "wuenlp_tools.visualize",
 29)
 30
 31# Fallback capabilities for steps that do not set provides/requires on PipelineStep.
 32STEP_CAPABILITIES: dict[str, dict[str, list[PipelineCapability]]] = {
 33    "SentimentLexiconBaseline": {"requires": [PipelineCapability.SCENES]},
 34    "LLM Sentiment": {"requires": [PipelineCapability.SCENES]},
 35    "HarryMotions Prompting": {"requires": [PipelineCapability.SCENES, PipelineCapability.CHARACTERS]},
 36    "HarryMotions Entity 2cl Prompting": {"requires": [PipelineCapability.SCENES, PipelineCapability.CHARACTERS]},
 37    "HarryMotions Entity 5cl Prompting": {"requires": [PipelineCapability.SCENES, PipelineCapability.CHARACTERS]},
 38    "HarryMotions Entity 8cl Prompting": {"requires": [PipelineCapability.SCENES, PipelineCapability.CHARACTERS]},
 39    "HarryMotions Aggregation": {"requires": [PipelineCapability.SCENES, PipelineCapability.CHARACTERS]},
 40    "Llama HarryMotions": {"requires": [PipelineCapability.PREPROCESS]},
 41    "PromptingDangerAnnotator": {"requires": [PipelineCapability.SCENES]},
 42    "PromptingFearAnnotator": {"requires": [PipelineCapability.SCENES]},
 43    "Wordlist Danger Annotator": {"requires": [PipelineCapability.SCENES]},
 44    "Wordlist Fear Annotator": {"requires": [PipelineCapability.SCENES]},
 45    "BERT Any Suspense": {"requires": [PipelineCapability.PREPROCESS, PipelineCapability.SCENES]},
 46    "BERT Dangerous Situation": {"requires": [PipelineCapability.PREPROCESS, PipelineCapability.SCENES]},
 47    "BERT Fear Description": {"requires": [PipelineCapability.PREPROCESS, PipelineCapability.SCENES]},
 48    "Suspense Aggregation": {"requires": [PipelineCapability.SCENES]},
 49    "SceneVisualiser": {"requires": [PipelineCapability.SCENES]},
 50}
 51
 52DEFAULT_CAPABILITY_PROVIDERS: dict[PipelineCapability, str] = {
 53    PipelineCapability.PREPROCESS: "wuenlp_tools.models.preprocess:Preprocessor",
 54    PipelineCapability.COREF: "wuenlp_tools.models.preprocess:Preprocessor",
 55    PipelineCapability.CHARACTERS: "wuenlp_tools.models.preprocess:Preprocessor",
 56    PipelineCapability.SCENES: "wuenlp_tools.models.scenes.segmentation:BERTSceneSegmenter",
 57    PipelineCapability.SCENE_SUMMARIES: "wuenlp_tools.models.scenes.summarisation:SceneSummarizer",
 58}
 59
 60_SKIP_MODULE_SUBSTRINGS = (
 61    ".tests.",
 62    ".test_",
 63)
 64
 65
 66@dataclass(frozen=True)
 67class StepCatalogEntry:
 68    id: str
 69    name: str
 70    category: str
 71    description: str
 72    module: str
 73    attr: str
 74    is_factory: bool
 75    requires_api_key: bool = False
 76    requires_paid: bool = False
 77    default_overwrite: bool = False
 78    provides: tuple[PipelineCapability, ...] = ()
 79    requires: tuple[PipelineCapability, ...] = ()
 80
 81
 82_CATALOG_CACHE: dict[str, StepCatalogEntry] | None = None
 83
 84
 85def _step_id(module_name: str, attr_name: str) -> str:
 86    return f"{module_name}:{attr_name}"
 87
 88
 89def _category_from_module(module_name: str) -> str:
 90    rel = module_name.removeprefix("wuenlp_tools.")
 91    parts = rel.split(".")
 92    if not rel:
 93        return "Other"
 94    if parts[0] == "models" and len(parts) > 1:
 95        return parts[1].replace("_", " ").title()
 96    if parts[0] == "visualize":
 97        return "Visualize"
 98    return parts[0].replace("_", " ").title()
 99
100
101def _is_step_instance(obj: Any) -> bool:
102    return isinstance(obj, PipelineStep) and not inspect.isclass(obj)
103
104
105def _is_step_factory(obj: Any) -> bool:
106    if not callable(obj) or inspect.isclass(obj):
107        return False
108    name = getattr(obj, "__name__", "")
109    if name == "PipelineStep" or name.startswith("_"):
110        return False
111    if not _STEP_FACTORY_NAME.match(name):
112        return False
113    try:
114        signature = inspect.signature(obj)
115    except (TypeError, ValueError):
116        return False
117    for param in signature.parameters.values():
118        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
119            continue
120        if param.default is inspect.Parameter.empty and param.name not in {"self", "cls"}:
121            return False
122    return True
123
124
125def _step_description(step: PipelineStep, module_name: str, attr_name: str) -> str:
126    processor_doc = inspect.getdoc(step.processor.__class__)
127    if processor_doc:
128        first_line = processor_doc.strip().splitlines()[0]
129        return first_line
130    return f"{module_name}.{attr_name}"
131
132
133def _capabilities_for_step(step: PipelineStep) -> tuple[list[PipelineCapability], list[PipelineCapability]]:
134    if step.provides or step.requires:
135        return list(step.provides), list(step.requires)
136    fallback = STEP_CAPABILITIES.get(step.name or "", {})
137    return fallback.get("provides", []), fallback.get("requires", [])
138
139
140def _entry_from_step(
141    module_name: str,
142    attr_name: str,
143    step: PipelineStep,
144    *,
145    is_factory: bool,
146) -> StepCatalogEntry:
147    provides, requires = _capabilities_for_step(step)
148    return StepCatalogEntry(
149        id=_step_id(module_name, attr_name),
150        name=step.name or attr_name,
151        category=_category_from_module(module_name),
152        description=_step_description(step, module_name, attr_name),
153        module=module_name,
154        attr=attr_name,
155        is_factory=is_factory,
156        requires_api_key=step.requires_api_key is not None,
157        requires_paid=bool(step.requires_paid_api_requests),
158        default_overwrite=step.overwrite is True,
159        provides=tuple(provides),
160        requires=tuple(requires),
161    )
162
163
164def _should_skip_module(module_name: str) -> bool:
165    if any(module_name == prefix or module_name.startswith(prefix + ".") for prefix in _SKIP_MODULE_PREFIXES):
166        return True
167    return any(part in module_name for part in _SKIP_MODULE_SUBSTRINGS)
168
169
170def _discover_module_steps(module_name: str) -> list[StepCatalogEntry]:
171    try:
172        module = importlib.import_module(module_name)
173    except Exception:
174        return []
175
176    instances: list[tuple[str, PipelineStep]] = []
177    factories: list[tuple[str, Any]] = []
178
179    for attr_name, obj in inspect.getmembers(module):
180        if attr_name.startswith("_"):
181            continue
182        if _is_step_instance(obj):
183            instances.append((attr_name, obj))
184        elif _is_step_factory(obj):
185            factories.append((attr_name, obj))
186
187    instance_step_names = {step.name for _, step in instances}
188    entries: list[StepCatalogEntry] = []
189
190    for attr_name, step in instances:
191        entries.append(_entry_from_step(module_name, attr_name, step, is_factory=False))
192
193    for attr_name, factory in factories:
194        try:
195            sample = factory()
196        except Exception:
197            continue
198        if not _is_step_instance(sample):
199            continue
200        if sample.name in instance_step_names:
201            continue
202        entries.append(_entry_from_step(module_name, attr_name, sample, is_factory=True))
203
204    return entries
205
206
207def discover_pipeline_steps() -> dict[str, StepCatalogEntry]:
208    raw_entries: list[StepCatalogEntry] = []
209    for root_name in _DISCOVERY_ROOTS:
210        root_pkg = importlib.import_module(root_name)
211        for modinfo in pkgutil.walk_packages(root_pkg.__path__, root_pkg.__name__ + "."):
212            if _should_skip_module(modinfo.name):
213                continue
214            raw_entries.extend(_discover_module_steps(modinfo.name))
215
216    entries: dict[str, StepCatalogEntry] = {}
217    seen_instance_ids: set[int] = set()
218    seen_factory_ids: set[int] = set()
219    instance_names: set[str] = set()
220
221    for entry in sorted(raw_entries, key=lambda item: (item.is_factory, item.id)):
222        if entry.is_factory:
223            if entry.name in instance_names:
224                continue
225            factory = getattr(importlib.import_module(entry.module), entry.attr)
226            factory_id = id(factory)
227            if factory_id in seen_factory_ids:
228                continue
229            seen_factory_ids.add(factory_id)
230        else:
231            step = _instantiate_step(entry, {})
232            instance_id = id(step)
233            if instance_id in seen_instance_ids:
234                continue
235            seen_instance_ids.add(instance_id)
236            instance_names.add(entry.name)
237
238        entries[entry.id] = entry
239
240    return dict(sorted(entries.items(), key=lambda item: (item[1].category, item[1].name)))
241
242
243def get_step_catalog() -> dict[str, StepCatalogEntry]:
244    global _CATALOG_CACHE
245    if _CATALOG_CACHE is None:
246        _CATALOG_CACHE = discover_pipeline_steps()
247    return _CATALOG_CACHE
248
249
250def clear_step_catalog_cache() -> None:
251    global _CATALOG_CACHE
252    _CATALOG_CACHE = None
253
254
255# Presets reference stable step display names (PipelineStep.name), not catalog ids.
256PIPELINE_PRESET_STEP_NAMES: dict[str, list[str]] = {
257    "empty": [],
258    "process_file": [
259        "Preprocess",
260        "BERT Scene segmentation",
261        "Scene Summarizer",
262        "LLM Sentiment",
263        "Main Character Extractor",
264    ],
265    "example": [
266        "Preprocess",
267        "Main Character Extractor",
268        "BERT Scene segmentation",
269        "Scene Summarizer",
270        "Scene Embedder",
271        "SentimentLexiconBaseline",
272        "Wordlist Danger Annotator",
273        "PromptingDangerAnnotator",
274        "LLM Sentiment",
275        "HarryMotions Prompting",
276        "MomentEmbedder",
277    ],
278    "coref_only": [
279        "Preprocess",
280        "LLM Alias Coref",
281    ],
282}
283
284
285def _resolve_step_ref(step_ref: str) -> StepCatalogEntry:
286    catalog = get_step_catalog()
287    if step_ref in catalog:
288        return catalog[step_ref]
289    matches = [entry for entry in catalog.values() if entry.name == step_ref]
290    if len(matches) == 1:
291        return matches[0]
292    if len(matches) > 1:
293        ids = ", ".join(entry.id for entry in matches)
294        raise KeyError(f"Ambiguous pipeline step name {step_ref!r}. Use one of: {ids}")
295    raise KeyError(f"Unknown pipeline step: {step_ref}")
296
297
298def get_pipeline_presets() -> dict[str, list[str]]:
299    return {
300        preset_name: order_pipeline_steps(
301            [_resolve_step_ref(step_name).id for step_name in step_names],
302        )
303        for preset_name, step_names in PIPELINE_PRESET_STEP_NAMES.items()
304    }
305
306
307def order_pipeline_steps(step_refs: list[str]) -> list[str]:
308    """Reorder selected steps so providers run before consumers. Does not add steps."""
309    catalog = get_step_catalog()
310    selected_ids = list(dict.fromkeys(_resolve_step_ref(step_ref).id for step_ref in step_refs))
311
312    all_provided: set[PipelineCapability] = set()
313    for step_id in selected_ids:
314        all_provided.update(catalog[step_id].provides)
315
316    missing: list[str] = []
317    for step_id in selected_ids:
318        entry = catalog[step_id]
319        for capability in entry.requires:
320            if capability not in all_provided:
321                missing.append(f"{entry.name!r} requires {capability.value!r}")
322    if missing:
323        providers = ", ".join(
324            f"{capability.value} (e.g. {DEFAULT_CAPABILITY_PROVIDERS[capability].split(':')[-1]})"
325            for capability in PipelineCapability
326            if capability in {c for sid in selected_ids for c in catalog[sid].requires}
327            and capability not in all_provided
328            and capability in DEFAULT_CAPABILITY_PROVIDERS
329        )
330        raise ValueError(
331            "Selected steps do not satisfy all dependencies: "
332            + "; ".join(missing)
333            + (f". Consider adding a step that provides: {providers}" if providers else "")
334        )
335
336    ordered: list[str] = []
337    provided: set[PipelineCapability] = set()
338    remaining = list(selected_ids)
339
340    while remaining:
341        ready = [
342            step_id for step_id in remaining
343            if all(capability in provided for capability in catalog[step_id].requires)
344        ]
345        if not ready:
346            blocked = remaining[0]
347            entry = catalog[blocked]
348            unsatisfied = [c for c in entry.requires if c not in provided]
349            raise ValueError(
350                f"Cannot order selected steps: {entry.name!r} requires "
351                f"{', '.join(c.value for c in unsatisfied)} before it runs, but no earlier "
352                f"selected step provides {'them' if len(unsatisfied) > 1 else 'it'}."
353            )
354        ready.sort(key=selected_ids.index)
355        for step_id in ready:
356            ordered.append(step_id)
357            provided.update(catalog[step_id].provides)
358            remaining.remove(step_id)
359
360    return ordered
361
362
363def catalog_as_dict() -> dict[str, Any]:
364    steps = []
365    for entry in get_step_catalog().values():
366        steps.append({
367            "id": entry.id,
368            "name": entry.name,
369            "category": entry.category,
370            "description": entry.description,
371            "requires_api_key": entry.requires_api_key,
372            "requires_paid": entry.requires_paid,
373            "default_overwrite": entry.default_overwrite,
374            "provides": [c.value for c in entry.provides],
375            "requires": [c.value for c in entry.requires],
376        })
377    capabilities = [
378        {
379            "id": capability.value,
380            "description": capability.description,
381            "affected_types": [t.__name__ for t in capability.affected_types],
382        }
383        for capability in PipelineCapability
384    ]
385    return {
386        "steps": steps,
387        "presets": get_pipeline_presets(),
388        "categories": list(dict.fromkeys(step["category"] for step in steps)),
389        "capabilities": capabilities,
390    }
391
392
393def catalog_as_json() -> str:
394    return json.dumps(catalog_as_dict(), ensure_ascii=False, indent=2)
395
396
397def _instantiate_step(entry: StepCatalogEntry, step_config: dict[str, Any]) -> PipelineStep:
398    module = importlib.import_module(entry.module)
399    obj = getattr(module, entry.attr)
400    if entry.is_factory:
401        if step_config:
402            return obj(**step_config)
403        return obj()
404    return obj
405
406
407def build_step(step_spec: str | dict[str, Any]) -> PipelineStep:
408    if isinstance(step_spec, str):
409        step_ref = step_spec
410        step_config: dict[str, Any] = {}
411        overwrite = None
412    else:
413        step_ref = step_spec["id"]
414        step_config = step_spec.get("config", {})
415        overwrite = step_spec.get("overwrite")
416
417    entry = _resolve_step_ref(step_ref)
418    step = _instantiate_step(entry, step_config)
419    if overwrite is True:
420        step.set_overwrite(True)
421    elif overwrite is False:
422        step.set_overwrite(False)
423    elif entry.default_overwrite:
424        step.set_overwrite(True)
425    return step
426
427
428def build_steps(step_specs: list[str | dict[str, Any]]) -> list[PipelineStep]:
429    plain_refs: list[str] = []
430    spec_by_id: dict[str, str | dict[str, Any]] = {}
431    for spec in step_specs:
432        if isinstance(spec, str):
433            step_id = _resolve_step_ref(spec).id
434            plain_refs.append(step_id)
435            spec_by_id.setdefault(step_id, spec)
436        else:
437            step_id = _resolve_step_ref(spec["id"]).id
438            plain_refs.append(step_id)
439            spec_by_id[step_id] = spec
440
441    ordered_ids = order_pipeline_steps(plain_refs)
442    return [build_step(spec_by_id.get(step_id, step_id)) for step_id in ordered_ids]