diff --git a/rampart/__init__.py b/rampart/__init__.py index 8a0f807..e80e7d7 100644 --- a/rampart/__init__.py +++ b/rampart/__init__.py @@ -6,7 +6,9 @@ Public API re-exports for convenient top-level access. """ -from rampart.attacks import Attacks +from importlib import import_module +from typing import TYPE_CHECKING + from rampart.core.adapter import AgentAdapter, Session from rampart.core.errors import DriverError, EvaluatorError, InfrastructureError from rampart.core.evaluator import BaseEvaluator, Evaluator @@ -41,11 +43,22 @@ ToolCall, Turn, ) -from rampart.drivers.llm import LLMDriver -from rampart.evaluators import LLMJudge, TranscriptScope -from rampart.probes import Probes from rampart.pytest_plugin._collection import record_result +if TYPE_CHECKING: + from rampart.attacks import Attacks + from rampart.drivers.llm import LLMDriver + from rampart.evaluators import LLMJudge, TranscriptScope + from rampart.probes import Probes + +__lazy_imports__: dict[str, tuple[str, str]] = { + "Attacks": ("rampart.attacks", "Attacks"), + "LLMDriver": ("rampart.drivers.llm", "LLMDriver"), + "LLMJudge": ("rampart.evaluators", "LLMJudge"), + "Probes": ("rampart.probes", "Probes"), + "TranscriptScope": ("rampart.evaluators", "TranscriptScope"), +} + __all__ = [ "AgentAdapter", "AppManifest", @@ -90,3 +103,31 @@ "resolve_as_attack", "resolve_as_probe", ] + + +def __getattr__(name: str) -> object: + """Load a PyRIT-backed public export only when requested. + + Args: + name: Name of the requested module attribute. + + Returns: + The requested public API object. + + Raises: + AttributeError: If ``name`` is not a lazy public export. + """ + try: + module_name, attribute_name = __lazy_imports__[name] + except KeyError: + message = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(message) from None + + value = getattr(import_module(module_name), attribute_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return eager and lazy public module attributes.""" + return sorted({*globals(), *__all__}) diff --git a/tests/unit/test_public_api.py b/tests/unit/test_public_api.py new file mode 100644 index 0000000..f22aba6 --- /dev/null +++ b/tests/unit/test_public_api.py @@ -0,0 +1,78 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the top-level RAMPART public API.""" + +from __future__ import annotations + +import subprocess # ruff: ignore[suspicious-subprocess-import] +import sys +from types import SimpleNamespace +from typing import Any + +import pytest + +import rampart + + +def test_pytest_plugin_import_does_not_load_heavy_dependencies() -> None: + """Plugin startup should not import heavy execution dependencies.""" + script = """ +import sys + +import rampart.pytest_plugin.plugin + +heavy_modules = sorted( + name + for name in sys.modules + if name == "pyrit" + or name.startswith("pyrit.") + or name == "transformers" + or name.startswith("transformers.") +) +if heavy_modules: + raise SystemExit(f"unexpected heavy imports: {heavy_modules}") +""" + + result = subprocess.run( # ruff: ignore[subprocess-without-shell-equals-true] + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize( + ("name", "module_name", "attribute_name"), + [ + ("Attacks", "rampart.attacks", "Attacks"), + ("LLMDriver", "rampart.drivers.llm", "LLMDriver"), + ("LLMJudge", "rampart.evaluators", "LLMJudge"), + ("Probes", "rampart.probes", "Probes"), + ("TranscriptScope", "rampart.evaluators", "TranscriptScope"), + ], +) +def test_heavy_public_export_is_loaded_on_demand( + monkeypatch: pytest.MonkeyPatch, + name: str, + module_name: str, + attribute_name: str, +) -> None: + """A deferred top-level export should resolve once and then be cached.""" + sentinel = object() + previous: Any = rampart.__dict__.pop(name, None) + + def fake_import_module(requested_module: str) -> SimpleNamespace: + assert requested_module == module_name + return SimpleNamespace(**{attribute_name: sentinel}) + + monkeypatch.setattr(rampart, "import_module", fake_import_module) + try: + assert getattr(rampart, name) is sentinel + assert rampart.__dict__[name] is sentinel + finally: + rampart.__dict__.pop(name, None) + if previous is not None: + rampart.__dict__[name] = previous