Skip to content

Registry

The registry module provides decorator-based plugin registration.

Use these decorators for new methods:

  • @register_probe("name")
  • @register_monitor("name")
  • @register_attributor("name")

The pipeline runner calls create_probe, create_monitor, and create_attributor to instantiate methods from YAML.

Minimal example:

from collections.abc import Sequence
from typing import Any

from SafeLens.core.base import BaseProbe, Batch, ModelWrapper, ProbeResult
from SafeLens.core.registry import create_probe, register_probe


@register_probe("constant_probe")
class ConstantProbe(BaseProbe):
    def attach(self, model: ModelWrapper, layers: Sequence[int]) -> None:
        pass

    def detect(self, batch: Batch) -> ProbeResult:
        return ProbeResult(risk_score=0.0, critical_layers=[])

    def intervene(self, batch: Batch, direction: Any, scale: float) -> None:
        pass

    def detach(self) -> None:
        pass


probe = create_probe("constant_probe")

Decorator-based registries for pluggable SafeLens methods.

RegistryError

Bases: KeyError

Raised when a method registry lookup or registration fails.

Source code in src/SafeLens/core/registry.py 
20
21
class RegistryError(KeyError):
    """Raised when a method registry lookup or registration fails."""

create_attributor(name, config=None)

Instantiate a registered attributor.

Source code in src/SafeLens/core/registry.py 
104
105
106
def create_attributor(name: str, config: dict[str, Any] | None = None) -> BaseAttributor:
    """Instantiate a registered attributor."""
    return get_attributor(name)(config=config)

create_monitor(name, config=None)

Instantiate a registered monitor.

Source code in src/SafeLens/core/registry.py 
 99
100
101
def create_monitor(name: str, config: dict[str, Any] | None = None) -> BaseMonitor:
    """Instantiate a registered monitor."""
    return get_monitor(name)(config=config)

create_probe(name, config=None)

Instantiate a registered probe.

Source code in src/SafeLens/core/registry.py 
94
95
96
def create_probe(name: str, config: dict[str, Any] | None = None) -> BaseProbe:
    """Instantiate a registered probe."""
    return get_probe(name)(config=config)

get_attributor(name)

Return a registered attributor class.

Source code in src/SafeLens/core/registry.py 
86
87
88
89
90
91
def get_attributor(name: str) -> type[BaseAttributor]:
    """Return a registered attributor class."""
    try:
        return _ATTRIBUTOR_REGISTRY[name]
    except KeyError as exc:
        raise RegistryError(_unknown_name_message("attributor", name, list_attributors())) from exc

get_monitor(name)

Return a registered monitor class.

Source code in src/SafeLens/core/registry.py 
78
79
80
81
82
83
def get_monitor(name: str) -> type[BaseMonitor]:
    """Return a registered monitor class."""
    try:
        return _MONITOR_REGISTRY[name]
    except KeyError as exc:
        raise RegistryError(_unknown_name_message("monitor", name, list_monitors())) from exc

get_probe(name)

Return a registered probe class.

Source code in src/SafeLens/core/registry.py 
70
71
72
73
74
75
def get_probe(name: str) -> type[BaseProbe]:
    """Return a registered probe class."""
    try:
        return _PROBE_REGISTRY[name]
    except KeyError as exc:
        raise RegistryError(_unknown_name_message("probe", name, list_probes())) from exc

list_attributors()

List registered attributor names.

Source code in src/SafeLens/core/registry.py 
119
120
121
def list_attributors() -> list[str]:
    """List registered attributor names."""
    return sorted(_ATTRIBUTOR_REGISTRY)

list_monitors()

List registered monitor names.

Source code in src/SafeLens/core/registry.py 
114
115
116
def list_monitors() -> list[str]:
    """List registered monitor names."""
    return sorted(_MONITOR_REGISTRY)

list_probes()

List registered probe names.

Source code in src/SafeLens/core/registry.py 
109
110
111
def list_probes() -> list[str]:
    """List registered probe names."""
    return sorted(_PROBE_REGISTRY)

load_builtin_methods()

Import built-in plugins so their registration decorators run.

Source code in src/SafeLens/core/registry.py 
24
25
26
27
28
29
30
31
def load_builtin_methods() -> None:
    """Import built-in plugins so their registration decorators run."""
    import SafeLens.attribution.captum  # noqa: F401
    import SafeLens.attribution.dummy  # noqa: F401
    import SafeLens.attribution.safety_heads  # noqa: F401
    import SafeLens.monitors.dummy  # noqa: F401
    import SafeLens.probes.dummy  # noqa: F401
    import SafeLens.probes.linear  # noqa: F401

register_attributor(name, replace=False)

Register an attributor class by name.

Source code in src/SafeLens/core/registry.py 
65
66
67
def register_attributor(name: str, replace: bool = False) -> Callable[[AttributorT], AttributorT]:
    """Register an attributor class by name."""
    return _register(_ATTRIBUTOR_REGISTRY, "attributor", name, replace)  # type: ignore[return-value]

register_monitor(name, replace=False)

Register a monitor class by name.

Source code in src/SafeLens/core/registry.py 
60
61
62
def register_monitor(name: str, replace: bool = False) -> Callable[[MonitorT], MonitorT]:
    """Register a monitor class by name."""
    return _register(_MONITOR_REGISTRY, "monitor", name, replace)  # type: ignore[return-value]

register_probe(name, replace=False)

Register a probe class by name.

Source code in src/SafeLens/core/registry.py 
55
56
57
def register_probe(name: str, replace: bool = False) -> Callable[[ProbeT], ProbeT]:
    """Register a probe class by name."""
    return _register(_PROBE_REGISTRY, "probe", name, replace)  # type: ignore[return-value]