diff --git a/eval_protocol/__init__.py b/eval_protocol/__init__.py index 5939896f..1ff39712 100644 --- a/eval_protocol/__init__.py +++ b/eval_protocol/__init__.py @@ -8,32 +8,17 @@ tool-augmented models using self-contained task bundles. """ +from importlib import import_module +from typing import Any import warnings -from eval_protocol.adapters.braintrust import reward_fn_to_scorer, scorer_to_reward_fn - +# Lightweight imports (no heavy optional dependencies) +from .integrations.braintrust import reward_fn_to_scorer, scorer_to_reward_fn 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, -) - -# Try to import FireworksPolicy if available -try: - from .mcp_env import FireworksPolicy - _FIREWORKS_AVAILABLE = True -except (ImportError, AttributeError): - _FIREWORKS_AVAILABLE = False # Import submodules to make them available via eval_protocol.rewards, etc. -from . import mcp, rewards from .models import EvaluateResult, Message, MetricResult from .playback_policy import PlaybackPolicyBase from .resources import create_llm_resource @@ -42,6 +27,7 @@ warnings.filterwarnings("default", category=DeprecationWarning, module="eval_protocol") +# Public API (static exports only; dynamic MCP symbols are provided via __getattr__) __all__ = [ # Core interfaces "Message", @@ -60,14 +46,6 @@ "RewardKitConfig", # Utilities "load_jsonl", - # MCP Environment API - "make", - "rollout", - "LiteLLMPolicy", - "AnthropicPolicy", - "FireworksPolicy", - "OpenAIPolicy", - "test_mcp", # Playback functionality "PlaybackPolicyBase", # Resource management @@ -77,6 +55,30 @@ "mcp", ] + +def __getattr__(name: str) -> Any: + """Lazily import heavy MCP environment symbols to speed up package import. + + This defers importing modules that depend on optional or heavy dependencies + (e.g., vendored tau2, OpenAI clients) until they are actually used. + """ + if name in { + "make", + "rollout", + "LiteLLMPolicy", + "AnthropicPolicy", + "FireworksPolicy", + "OpenAIPolicy", + "test_mcp", + }: + m = import_module(".mcp_env", __name__) + return getattr(m, name) + if name in {"mcp", "rewards"}: + # Lazy-load subpackages for attribute access like eval_protocol.mcp + return import_module(f".{name}", __name__) + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + from . import _version __version__ = _version.get_versions()["version"] diff --git a/eval_protocol/adapters/__init__.py b/eval_protocol/adapters/__init__.py index fc04237b..efbbeac1 100644 --- a/eval_protocol/adapters/__init__.py +++ b/eval_protocol/adapters/__init__.py @@ -10,38 +10,24 @@ - TRL integration (legacy) """ -# Conditional imports based on available dependencies -try: - from .langfuse import LangfuseAdapter, create_langfuse_adapter - __all__ = ["LangfuseAdapter", "create_langfuse_adapter"] -except ImportError: - __all__ = [] +from importlib import import_module +from typing import Any -try: - from .huggingface import ( - HuggingFaceAdapter, - create_huggingface_adapter, - create_gsm8k_adapter, - create_math_adapter, - ) - __all__.extend([ - "HuggingFaceAdapter", - "create_huggingface_adapter", - "create_gsm8k_adapter", - "create_math_adapter", - ]) -except ImportError: - pass +__all__ = [] -# Legacy adapters (always available) -try: - from .braintrust import reward_fn_to_scorer, scorer_to_reward_fn - __all__.extend(["scorer_to_reward_fn", "reward_fn_to_scorer"]) -except ImportError: - pass -try: - from .trl import create_trl_adapter - __all__.extend(["create_trl_adapter"]) -except ImportError: - pass +def __getattr__(name: str) -> Any: + # Lazy import optional adapters to avoid import-time side effects and heavy deps + if name in {"LangfuseAdapter", "create_langfuse_adapter"}: + m = import_module(".langfuse", __name__) + return getattr(m, name) + if name in {"HuggingFaceAdapter", "create_huggingface_adapter", "create_gsm8k_adapter", "create_math_adapter"}: + m = import_module(".huggingface", __name__) + return getattr(m, name) + if name in {"reward_fn_to_scorer", "scorer_to_reward_fn"}: + m = import_module(".braintrust", __name__) + return getattr(m, name) + if name in {"create_trl_adapter"}: + m = import_module(".trl", __name__) + return getattr(m, name) + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/vendor/tau2/utils/llm_utils.py b/vendor/tau2/utils/llm_utils.py index 98ae698c..244d11a4 100644 --- a/vendor/tau2/utils/llm_utils.py +++ b/vendor/tau2/utils/llm_utils.py @@ -7,6 +7,7 @@ from litellm.caching.caching import Cache from litellm.main import ModelResponse, Usage from loguru import logger +import os from vendor.tau2.config import ( DEFAULT_LLM_CACHE_TYPE, @@ -70,9 +71,8 @@ ALLOW_SONNET_THINKING = False - -if not ALLOW_SONNET_THINKING: - logger.warning("Sonnet thinking is disabled") +if os.getenv("TAU2_VERBOSE") == "1" and not ALLOW_SONNET_THINKING: + logger.info("Sonnet thinking is disabled") def _parse_ft_model_name(model: str) -> str: diff --git a/vendor/tau2/utils/utils.py b/vendor/tau2/utils/utils.py index c1103fe9..6c009501 100644 --- a/vendor/tau2/utils/utils.py +++ b/vendor/tau2/utils/utils.py @@ -9,8 +9,10 @@ from dotenv import load_dotenv from loguru import logger +_TAU2_VERBOSE = os.getenv("TAU2_VERBOSE") == "1" + res = load_dotenv() -if not res: +if not res and _TAU2_VERBOSE: logger.warning("No .env file found") # Try to get data directory from environment variable first @@ -19,15 +21,17 @@ if DATA_DIR_ENV: # Use environment variable if set DATA_DIR = Path(DATA_DIR_ENV) - logger.info(f"Using data directory from environment: {DATA_DIR}") + if _TAU2_VERBOSE: + logger.info(f"Using data directory from environment: {DATA_DIR}") else: # Fallback to vendored tau2 directory SOURCE_DIR = Path(__file__).parents[1] # vendor/tau2/ DATA_DIR = SOURCE_DIR / "data" - logger.info(f"Using data directory from vendored tau2: {DATA_DIR}") + if _TAU2_VERBOSE: + logger.info(f"Using data directory from vendored tau2: {DATA_DIR}") # Check if data directory exists and is accessible -if not DATA_DIR.exists(): +if not DATA_DIR.exists() and _TAU2_VERBOSE: logger.warning(f"Data directory does not exist: {DATA_DIR}") logger.warning( "Set TAU2_DATA_DIR environment variable to point to your data directory"