lazy loading of all modules in init

eval_protocol/__init__.py

@@ -8,102 +8,159 @@
 tool-augmented models using self-contained task bundles.
 """
 
+import importlib
+import sys
 import warnings
+from typing import TYPE_CHECKING
 
-from .auth import get_fireworks_account_id, get_fireworks_api_key
-from .common_utils import load_jsonl
-from .config import RewardKitConfig, get_config, load_config
-from .mcp_env import (
-    AnthropicPolicy,
-    FireworksPolicy,
-    LiteLLMPolicy,
-    OpenAIPolicy,
-    make,
-    rollout,
-    test_mcp,
-)
-from .data_loader import DynamicDataLoader, InlineDataLoader
-from . import mcp, rewards
-from .models import EvaluateResult, Message, MetricResult, EvaluationRow, InputMetadata, Status
-from .playback_policy import PlaybackPolicyBase
-from .resources import create_llm_resource
-from .reward_function import RewardFunction
-from .typed_interface import reward_function
-from .quickstart.aha_judge import aha_judge
-from .utils.evaluation_row_utils import (
-    multi_turn_assistant_to_ground_truth,
-    assistant_to_ground_truth,
-    filter_longest_conversation,
-)
-from .pytest import evaluation_test, SingleTurnRolloutProcessor, RemoteRolloutProcessor, GithubActionRolloutProcessor
-from .pytest.parameterize import DefaultParameterIdGenerator
-from .log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
-from .log_utils.rollout_id_filter import RolloutIdFilter
-from .log_utils.util import setup_rollout_logging_for_elasticsearch_handler
-from .log_utils.fireworks_tracing_http_handler import FireworksTracingHttpHandler
-from .log_utils.elasticsearch_client import ElasticsearchConfig
-
-
-from .types.remote_rollout_processor import (
-    InitRequest,
-    RolloutMetadata,
-    StatusResponse,
-    create_langfuse_config_tags,
-    DataLoaderConfig,
-)
-
-try:
-    from .adapters import OpenAIResponsesAdapter
-except ImportError:
-    OpenAIResponsesAdapter = None
-
-try:
-    from .adapters import LangfuseAdapter, create_langfuse_adapter
-except ImportError:
-    LangfuseAdapter = None
-
-try:
-    from .adapters import BraintrustAdapter, create_braintrust_adapter
-except ImportError:
-    BraintrustAdapter = None
-
-try:
-    from .adapters import LangSmithAdapter
-except ImportError:
-    LangSmithAdapter = None
-
-
-try:
-    from .adapters import WeaveAdapter
-except ImportError:
-    WeaveAdapter = None
-
-try:
-    from .proxy import create_app, AuthProvider, AccountInfo  # pyright: ignore[reportAssignmentType]
-except ImportError:
-
-    def create_app(*args, **kwargs):
-        raise ImportError(
-            "Proxy functionality requires additional dependencies. "
-            "Please install with: pip install eval-protocol[proxy]"
-        )
-
-    class AuthProvider:
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                "Proxy functionality requires additional dependencies. "
-                "Please install with: pip install eval-protocol[proxy]"
-            )
-
-    class AccountInfo:
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                "Proxy functionality requires additional dependencies. "
-                "Please install with: pip install eval-protocol[proxy]"
-            )
+warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol")
 
+# Eager imports for symbols that conflict with module names - ONLY when pytest is running.
+# The reward_function.py module exports RewardFunction class, and we also export the
+# reward_function decorator from typed_interface. When pytest's AssertionRewritingHook
+# imports eval_protocol.reward_function as a module, Python would set
+# eval_protocol.reward_function to the module, shadowing our function export.
+#
+# We detect pytest by checking if _pytest or pytest is already loaded. This avoids
+# the ~500ms import overhead for non-test scenarios like the CLI.
+_running_under_pytest = "_pytest" in sys.modules or "pytest" in sys.modules
+if _running_under_pytest:
+    from .reward_function import RewardFunction  # noqa: E402
+    from .typed_interface import reward_function  # noqa: E402
+
+# Lazy import mappings: name -> (module_path, attribute_name or None for module import)
+_LAZY_IMPORTS = {
+    # From .auth
+    "get_fireworks_account_id": (".auth", "get_fireworks_account_id"),
+    "get_fireworks_api_key": (".auth", "get_fireworks_api_key"),
+    # From .common_utils
+    "load_jsonl": (".common_utils", "load_jsonl"),
+    # From .config
+    "RewardKitConfig": (".config", "RewardKitConfig"),
+    "get_config": (".config", "get_config"),
+    "load_config": (".config", "load_config"),
+    # From .mcp_env
+    "AnthropicPolicy": (".mcp_env", "AnthropicPolicy"),
+    "FireworksPolicy": (".mcp_env", "FireworksPolicy"),
+    "LiteLLMPolicy": (".mcp_env", "LiteLLMPolicy"),
+    "OpenAIPolicy": (".mcp_env", "OpenAIPolicy"),
+    "make": (".mcp_env", "make"),
+    "rollout": (".mcp_env", "rollout"),
+    "test_mcp": (".mcp_env", "test_mcp"),
+    # From .data_loader
+    "DynamicDataLoader": (".data_loader", "DynamicDataLoader"),
+    "InlineDataLoader": (".data_loader", "InlineDataLoader"),
+    # Submodules (accessible as eval_protocol.submodule)
+    "mcp": (".mcp", None),
+    "rewards": (".rewards", None),
+    "models": (".models", None),
+    "auth": (".auth", None),
+    "config": (".config", None),
+    # From .models
+    "EvaluateResult": (".models", "EvaluateResult"),
+    "Message": (".models", "Message"),
+    "MetricResult": (".models", "MetricResult"),
+    "EvaluationRow": (".models", "EvaluationRow"),
+    "InputMetadata": (".models", "InputMetadata"),
+    "Status": (".models", "Status"),
+    # From .playback_policy
+    "PlaybackPolicyBase": (".playback_policy", "PlaybackPolicyBase"),
+    # From .resources
+    "create_llm_resource": (".resources", "create_llm_resource"),
+    # From .reward_function
+    "RewardFunction": (".reward_function", "RewardFunction"),
+    # From .typed_interface
+    "reward_function": (".typed_interface", "reward_function"),
+    # From .quickstart.aha_judge
+    "aha_judge": (".quickstart.aha_judge", "aha_judge"),
+    # From .utils.evaluation_row_utils
+    "multi_turn_assistant_to_ground_truth": (".utils.evaluation_row_utils", "multi_turn_assistant_to_ground_truth"),
+    "assistant_to_ground_truth": (".utils.evaluation_row_utils", "assistant_to_ground_truth"),
+    "filter_longest_conversation": (".utils.evaluation_row_utils", "filter_longest_conversation"),
+    # From .pytest
+    "evaluation_test": (".pytest", "evaluation_test"),
+    "SingleTurnRolloutProcessor": (".pytest", "SingleTurnRolloutProcessor"),
+    "RemoteRolloutProcessor": (".pytest", "RemoteRolloutProcessor"),
+    "GithubActionRolloutProcessor": (".pytest", "GithubActionRolloutProcessor"),
+    # From .pytest.parameterize
+    "DefaultParameterIdGenerator": (".pytest.parameterize", "DefaultParameterIdGenerator"),
+    # From .log_utils
+    "ElasticsearchDirectHttpHandler": (
+        ".log_utils.elasticsearch_direct_http_handler",
+        "ElasticsearchDirectHttpHandler",
+    ),
+    "RolloutIdFilter": (".log_utils.rollout_id_filter", "RolloutIdFilter"),
+    "setup_rollout_logging_for_elasticsearch_handler": (
+        ".log_utils.util",
+        "setup_rollout_logging_for_elasticsearch_handler",
+    ),
+    "FireworksTracingHttpHandler": (".log_utils.fireworks_tracing_http_handler", "FireworksTracingHttpHandler"),
+    "ElasticsearchConfig": (".log_utils.elasticsearch_client", "ElasticsearchConfig"),
+    # From .types.remote_rollout_processor
+    "InitRequest": (".types.remote_rollout_processor", "InitRequest"),
+    "RolloutMetadata": (".types.remote_rollout_processor", "RolloutMetadata"),
+    "StatusResponse": (".types.remote_rollout_processor", "StatusResponse"),
+    "create_langfuse_config_tags": (".types.remote_rollout_processor", "create_langfuse_config_tags"),
+    "DataLoaderConfig": (".types.remote_rollout_processor", "DataLoaderConfig"),
+}
+
+# Optional imports that may not be available
+_OPTIONAL_IMPORTS = {
+    "OpenAIResponsesAdapter": (".adapters", "OpenAIResponsesAdapter"),
+    "LangfuseAdapter": (".adapters", "LangfuseAdapter"),
+    "create_langfuse_adapter": (".adapters", "create_langfuse_adapter"),
+    "BraintrustAdapter": (".adapters", "BraintrustAdapter"),
+    "create_braintrust_adapter": (".adapters", "create_braintrust_adapter"),
+    "LangSmithAdapter": (".adapters", "LangSmithAdapter"),
+    "WeaveAdapter": (".adapters", "WeaveAdapter"),
+    "create_app": (".proxy", "create_app"),
+    "AuthProvider": (".proxy", "AuthProvider"),
+    "AccountInfo": (".proxy", "AccountInfo"),
+}
+
+
+def __getattr__(name: str):
+    """Lazy import handler for module-level attributes."""
+    # Check regular lazy imports
+    if name in _LAZY_IMPORTS:
+        module_path, attr_name = _LAZY_IMPORTS[name]
+        module = importlib.import_module(module_path, package="eval_protocol")
+        if attr_name is None:
+            # It's a submodule import
+            return module
+        return getattr(module, attr_name)
+
+    # Check optional imports
+    if name in _OPTIONAL_IMPORTS:
+        module_path, attr_name = _OPTIONAL_IMPORTS[name]
+        try:
+            module = importlib.import_module(module_path, package="eval_protocol")
+            return getattr(module, attr_name)
+        except ImportError:
+            # Return None or a placeholder for optional imports
+            if name in ("create_app",):
+
+                def create_app(*args, **kwargs):
+                    raise ImportError(
+                        "Proxy functionality requires additional dependencies. "
+                        "Please install with: pip install eval-protocol[proxy]"
+                    )
+
+                return create_app
+            elif name in ("AuthProvider", "AccountInfo"):
+
+                class _Placeholder:
+                    def __init__(self, *args, **kwargs):
+                        raise ImportError(
+                            "Proxy functionality requires additional dependencies. "
+                            "Please install with: pip install eval-protocol[proxy]"
+                        )
+
+                return _Placeholder
+            return None
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
-warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol")
 
 __all__ = [
     "ElasticsearchConfig",
@@ -173,6 +230,85 @@ def __init__(self, *args, **kwargs):
     "AccountInfo",
 ]
 
-from . import _version
+# Version is loaded lazily too
+_version_info = None
+
+
+def _get_version():
+    global _version_info
+    if _version_info is None:
+        from . import _version
+
+        _version_info = _version.get_versions()["version"]
+    return _version_info
+
+
+# For TYPE_CHECKING, we provide type hints so IDEs can see the exports
+if TYPE_CHECKING:
+    from .auth import get_fireworks_account_id, get_fireworks_api_key
+    from .common_utils import load_jsonl
+    from .config import RewardKitConfig, get_config, load_config
+    from .mcp_env import (
+        AnthropicPolicy,
+        FireworksPolicy,
+        LiteLLMPolicy,
+        OpenAIPolicy,
+        make,
+        rollout,
+        test_mcp,
+    )
+    from .data_loader import DynamicDataLoader, InlineDataLoader
+    from . import mcp, rewards
+    from .models import EvaluateResult, Message, MetricResult, EvaluationRow, InputMetadata, Status
+    from .playback_policy import PlaybackPolicyBase
+    from .resources import create_llm_resource
+    from .reward_function import RewardFunction
+    from .typed_interface import reward_function
+    from .quickstart.aha_judge import aha_judge
+    from .utils.evaluation_row_utils import (
+        multi_turn_assistant_to_ground_truth,
+        assistant_to_ground_truth,
+        filter_longest_conversation,
+    )
+    from .pytest import (
+        evaluation_test,
+        SingleTurnRolloutProcessor,
+        RemoteRolloutProcessor,
+        GithubActionRolloutProcessor,
+    )
+    from .pytest.parameterize import DefaultParameterIdGenerator
+    from .log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
+    from .log_utils.rollout_id_filter import RolloutIdFilter
+    from .log_utils.util import setup_rollout_logging_for_elasticsearch_handler
+    from .log_utils.fireworks_tracing_http_handler import FireworksTracingHttpHandler
+    from .log_utils.elasticsearch_client import ElasticsearchConfig
+    from .types.remote_rollout_processor import (
+        InitRequest,
+        RolloutMetadata,
+        StatusResponse,
+        create_langfuse_config_tags,
+        DataLoaderConfig,
+    )
+    from .adapters import (
+        OpenAIResponsesAdapter,
+        LangfuseAdapter,
+        create_langfuse_adapter,
+        BraintrustAdapter,
+        create_braintrust_adapter,
+        LangSmithAdapter,
+        WeaveAdapter,
+    )
+    from .proxy import create_app, AuthProvider, AccountInfo
+
+
+# __version__ property - computed lazily
+class _VersionModule:
+    @property
+    def __version__(self):
+        return _get_version()
+
+
+import sys
 
-__version__ = _version.get_versions()["version"]
+_this_module = sys.modules[__name__]
+_this_module.__class__ = type("module", (type(_this_module),), {"__version__": property(lambda self: _get_version())})

pytest.ini

@@ -1,6 +1,7 @@
 [pytest]
 markers =
     asyncio
+    benchmark: marks tests as benchmark tests (for CLI startup time checks)
 asyncio_mode = auto
 asyncio_default_fixture_loop_scope = function
 testpaths = tests ./eval_protocol/quickstart