From f1a216cd31b0bf858b5aff36d13d764c103068f3 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 22 Sep 2026 16:39:09 -0400 Subject: [PATCH 1/2] fix: Lazy load integration modules --- py/src/braintrust/__init__.py | 58 +++++++-- py/src/braintrust/auto.py | 100 ++++++--------- py/src/braintrust/integrations/__init__.py | 114 +++++++++++++----- .../test_lazy_integration_imports.py | 84 +++++++++++++ 4 files changed, 255 insertions(+), 101 deletions(-) create mode 100644 py/src/braintrust/test_lazy_integration_imports.py diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 7af51764..93a885fb 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -50,6 +50,7 @@ def is_equal(expected, output): # Check env var at import time for auto-instrumentation import os +from typing import Any as _Any if os.getenv("BRAINTRUST_INSTRUMENT_THREADS", "").lower() in ("true", "1", "yes"): @@ -68,13 +69,6 @@ def is_equal(expected, output): from .functions.invoke import * from .functions.stream import * from .generated_types import * -from .integrations.ai_sdk import setup_ai_sdk as setup_ai_sdk -from .integrations.anthropic import wrap_anthropic as wrap_anthropic -from .integrations.instructor import wrap_instructor as wrap_instructor -from .integrations.litellm import wrap_litellm as wrap_litellm -from .integrations.openai import wrap_openai as wrap_openai -from .integrations.openrouter import wrap_openrouter as wrap_openrouter -from .integrations.pydantic_ai import setup_pydantic_ai as setup_pydantic_ai from .logger import * from .logger import ( _internal_get_global_state, # noqa: F401 # type: ignore[reportUnusedImport] @@ -91,3 +85,53 @@ def is_equal(expected, output): from .span_customizer import set_span_customizers as set_span_customizers from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE from .util import MarkAsyncWrapper as MarkAsyncWrapper + + +def wrap_anthropic(client: _Any) -> _Any: + from .integrations.anthropic import wrap_anthropic as _wrap_anthropic + + return _wrap_anthropic(client) + + +def wrap_instructor(client: _Any) -> _Any: + from .integrations.instructor import wrap_instructor as _wrap_instructor + + return _wrap_instructor(client) + + +def wrap_litellm(litellm: _Any) -> _Any: + from .integrations.litellm import wrap_litellm as _wrap_litellm + + return _wrap_litellm(litellm) + + +def wrap_openai(client: _Any) -> _Any: + from .integrations.openai import wrap_openai as _wrap_openai + + return _wrap_openai(client) + + +def wrap_openrouter(client: _Any) -> _Any: + from .integrations.openrouter import wrap_openrouter as _wrap_openrouter + + return _wrap_openrouter(client) + + +def setup_ai_sdk( + api_key: str | None = None, + project_id: str | None = None, + project_name: str | None = None, +) -> bool: + from .integrations.ai_sdk import setup_ai_sdk as _setup_ai_sdk + + return _setup_ai_sdk(api_key=api_key, project_id=project_id, project_name=project_name) + + +def setup_pydantic_ai( + api_key: str | None = None, + project_id: str | None = None, + project_name: str | None = None, +) -> bool: + from .integrations.pydantic_ai import setup_pydantic_ai as _setup_pydantic_ai + + return _setup_pydantic_ai(api_key=api_key, project_id=project_id, project_name=project_name) diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index d0721813..084cca3d 100644 --- a/py/src/braintrust/auto.py +++ b/py/src/braintrust/auto.py @@ -7,40 +7,9 @@ import logging from collections.abc import Sequence from contextlib import contextmanager +from importlib import import_module +from typing import Any -from braintrust.integrations import ( - ADKIntegration, - AgentScopeIntegration, - AgnoIntegration, - AISDKIntegration, - AnthropicIntegration, - AutoGenIntegration, - BedrockRuntimeIntegration, - ClaudeAgentSDKIntegration, - CohereIntegration, - CrewAIIntegration, - CursorSDKIntegration, - DSPyIntegration, - GoogleDiscoveryEngineIntegration, - GoogleGenAIIntegration, - HuggingFaceHubIntegration, - InstructorIntegration, - LangChainIntegration, - LiteLLMIntegration, - LiveKitAgentsIntegration, - LlamaIndexIntegration, - MistralIntegration, - OpenAIAgentsIntegration, - OpenAIIntegration, - OpenRouterIntegration, - PipecatIntegration, - PydanticAIIntegration, - StrandsIntegration, - TemporalIntegration, - TransformersIntegration, - TypeSafeIntegration, -) -from braintrust.integrations.base import BaseIntegration from braintrust.span_customizer import SpanCustomizer, set_span_customizers @@ -188,70 +157,73 @@ def auto_instrument( results: dict[str, bool] = {} if openai: - results["openai"] = _instrument_integration(OpenAIIntegration) + results["openai"] = _instrument_integration("openai", "OpenAIIntegration") if anthropic: - results["anthropic"] = _instrument_integration(AnthropicIntegration) + results["anthropic"] = _instrument_integration("anthropic", "AnthropicIntegration") if litellm: - results["litellm"] = _instrument_integration(LiteLLMIntegration) + results["litellm"] = _instrument_integration("litellm", "LiteLLMIntegration") if ai_sdk: - results["ai_sdk"] = _instrument_integration(AISDKIntegration) + results["ai_sdk"] = _instrument_integration("ai_sdk", "AISDKIntegration") if pydantic_ai: - results["pydantic_ai"] = _instrument_integration(PydanticAIIntegration) + results["pydantic_ai"] = _instrument_integration("pydantic_ai", "PydanticAIIntegration") if google_genai: - results["google_genai"] = _instrument_integration(GoogleGenAIIntegration) + results["google_genai"] = _instrument_integration("google_genai", "GoogleGenAIIntegration") if google_discoveryengine: - results["google_discoveryengine"] = _instrument_integration(GoogleDiscoveryEngineIntegration) + results["google_discoveryengine"] = _instrument_integration( + "google_discoveryengine", "GoogleDiscoveryEngineIntegration" + ) if instructor: - results["instructor"] = _instrument_integration(InstructorIntegration) + results["instructor"] = _instrument_integration("instructor", "InstructorIntegration") if openrouter: - results["openrouter"] = _instrument_integration(OpenRouterIntegration) + results["openrouter"] = _instrument_integration("openrouter", "OpenRouterIntegration") if mistral: - results["mistral"] = _instrument_integration(MistralIntegration) + results["mistral"] = _instrument_integration("mistral", "MistralIntegration") if huggingface_hub: - results["huggingface_hub"] = _instrument_integration(HuggingFaceHubIntegration) + results["huggingface_hub"] = _instrument_integration("huggingface_hub", "HuggingFaceHubIntegration") if transformers: - results["transformers"] = _instrument_integration(TransformersIntegration) + results["transformers"] = _instrument_integration("transformers", "TransformersIntegration") if agno: - results["agno"] = _instrument_integration(AgnoIntegration) + results["agno"] = _instrument_integration("agno", "AgnoIntegration") if agentscope: - results["agentscope"] = _instrument_integration(AgentScopeIntegration) + results["agentscope"] = _instrument_integration("agentscope", "AgentScopeIntegration") if claude_agent_sdk: - results["claude_agent_sdk"] = _instrument_integration(ClaudeAgentSDKIntegration) + results["claude_agent_sdk"] = _instrument_integration("claude_agent_sdk", "ClaudeAgentSDKIntegration") if cursor_sdk: - results["cursor_sdk"] = _instrument_integration(CursorSDKIntegration) + results["cursor_sdk"] = _instrument_integration("cursor_sdk", "CursorSDKIntegration") if dspy: - results["dspy"] = _instrument_integration(DSPyIntegration) + results["dspy"] = _instrument_integration("dspy", "DSPyIntegration") if adk: - results["adk"] = _instrument_integration(ADKIntegration) + results["adk"] = _instrument_integration("adk", "ADKIntegration") if langchain: - results["langchain"] = _instrument_integration(LangChainIntegration) + results["langchain"] = _instrument_integration("langchain", "LangChainIntegration") if llamaindex: - results["llamaindex"] = _instrument_integration(LlamaIndexIntegration) + results["llamaindex"] = _instrument_integration("llamaindex", "LlamaIndexIntegration") if openai_agents: - results["openai_agents"] = _instrument_integration(OpenAIAgentsIntegration) + results["openai_agents"] = _instrument_integration("openai_agents", "OpenAIAgentsIntegration") if cohere: - results["cohere"] = _instrument_integration(CohereIntegration) + results["cohere"] = _instrument_integration("cohere", "CohereIntegration") if autogen: - results["autogen"] = _instrument_integration(AutoGenIntegration) + results["autogen"] = _instrument_integration("autogen", "AutoGenIntegration") if bedrock: - results["bedrock_runtime"] = _instrument_integration(BedrockRuntimeIntegration) + results["bedrock_runtime"] = _instrument_integration("bedrock_runtime", "BedrockRuntimeIntegration") if crewai: - results["crewai"] = _instrument_integration(CrewAIIntegration) + results["crewai"] = _instrument_integration("crewai", "CrewAIIntegration") if strands: - results["strands"] = _instrument_integration(StrandsIntegration) + results["strands"] = _instrument_integration("strands", "StrandsIntegration") if temporal: - results["temporal"] = _instrument_integration(TemporalIntegration) + results["temporal"] = _instrument_integration("temporal", "TemporalIntegration") if livekit_agents: - results["livekit_agents"] = _instrument_integration(LiveKitAgentsIntegration) + results["livekit_agents"] = _instrument_integration("livekit_agents", "LiveKitAgentsIntegration") if pipecat: - results["pipecat"] = _instrument_integration(PipecatIntegration) + results["pipecat"] = _instrument_integration("pipecat", "PipecatIntegration") if typesafe: - results["typesafe"] = _instrument_integration(TypeSafeIntegration) + results["typesafe"] = _instrument_integration("typesafe", "TypeSafeIntegration") return results -def _instrument_integration(integration: type[BaseIntegration]) -> bool: +def _instrument_integration(module_name: str, class_name: str) -> bool: with _try_patch(): + integration: Any = getattr(import_module(f"braintrust.integrations.{module_name}"), class_name) return integration.setup() return False diff --git a/py/src/braintrust/integrations/__init__.py b/py/src/braintrust/integrations/__init__.py index c391e7dc..38acb690 100644 --- a/py/src/braintrust/integrations/__init__.py +++ b/py/src/braintrust/integrations/__init__.py @@ -1,33 +1,74 @@ -from .adk import ADKIntegration -from .agentscope import AgentScopeIntegration -from .agno import AgnoIntegration -from .ai_sdk import AISDKIntegration -from .anthropic import AnthropicIntegration -from .autogen import AutoGenIntegration -from .bedrock_runtime import BedrockRuntimeIntegration -from .claude_agent_sdk import ClaudeAgentSDKIntegration -from .cohere import CohereIntegration -from .crewai import CrewAIIntegration -from .cursor_sdk import CursorSDKIntegration -from .dspy import DSPyIntegration -from .google_discoveryengine import GoogleDiscoveryEngineIntegration -from .google_genai import GoogleGenAIIntegration -from .huggingface_hub import HuggingFaceHubIntegration -from .instructor import InstructorIntegration -from .langchain import LangChainIntegration -from .litellm import LiteLLMIntegration -from .livekit_agents import LiveKitAgentsIntegration -from .llamaindex import LlamaIndexIntegration -from .mistral import MistralIntegration -from .openai import OpenAIIntegration -from .openai_agents import OpenAIAgentsIntegration -from .openrouter import OpenRouterIntegration -from .pipecat import PipecatIntegration -from .pydantic_ai import PydanticAIIntegration -from .strands import StrandsIntegration -from .temporal import TemporalIntegration -from .transformers import TransformersIntegration -from .typesafe import TypeSafeIntegration +"""Public integration classes, loaded only when their attributes are accessed.""" + +from importlib import import_module +from typing import TYPE_CHECKING, Any + + +_INTEGRATION_MODULES = { + "ADKIntegration": "adk", + "AgentScopeIntegration": "agentscope", + "AgnoIntegration": "agno", + "AISDKIntegration": "ai_sdk", + "AnthropicIntegration": "anthropic", + "AutoGenIntegration": "autogen", + "BedrockRuntimeIntegration": "bedrock_runtime", + "ClaudeAgentSDKIntegration": "claude_agent_sdk", + "CohereIntegration": "cohere", + "CrewAIIntegration": "crewai", + "CursorSDKIntegration": "cursor_sdk", + "DSPyIntegration": "dspy", + "GoogleDiscoveryEngineIntegration": "google_discoveryengine", + "GoogleGenAIIntegration": "google_genai", + "HuggingFaceHubIntegration": "huggingface_hub", + "InstructorIntegration": "instructor", + "LangChainIntegration": "langchain", + "LiteLLMIntegration": "litellm", + "LiveKitAgentsIntegration": "livekit_agents", + "LlamaIndexIntegration": "llamaindex", + "MistralIntegration": "mistral", + "OpenAIIntegration": "openai", + "OpenAIAgentsIntegration": "openai_agents", + "OpenRouterIntegration": "openrouter", + "PipecatIntegration": "pipecat", + "PydanticAIIntegration": "pydantic_ai", + "StrandsIntegration": "strands", + "TemporalIntegration": "temporal", + "TransformersIntegration": "transformers", + "TypeSafeIntegration": "typesafe", +} + + +if TYPE_CHECKING: + from .adk import ADKIntegration as ADKIntegration + from .agentscope import AgentScopeIntegration as AgentScopeIntegration + from .agno import AgnoIntegration as AgnoIntegration + from .ai_sdk import AISDKIntegration as AISDKIntegration + from .anthropic import AnthropicIntegration as AnthropicIntegration + from .autogen import AutoGenIntegration as AutoGenIntegration + from .bedrock_runtime import BedrockRuntimeIntegration as BedrockRuntimeIntegration + from .claude_agent_sdk import ClaudeAgentSDKIntegration as ClaudeAgentSDKIntegration + from .cohere import CohereIntegration as CohereIntegration + from .crewai import CrewAIIntegration as CrewAIIntegration + from .cursor_sdk import CursorSDKIntegration as CursorSDKIntegration + from .dspy import DSPyIntegration as DSPyIntegration + from .google_discoveryengine import GoogleDiscoveryEngineIntegration as GoogleDiscoveryEngineIntegration + from .google_genai import GoogleGenAIIntegration as GoogleGenAIIntegration + from .huggingface_hub import HuggingFaceHubIntegration as HuggingFaceHubIntegration + from .instructor import InstructorIntegration as InstructorIntegration + from .langchain import LangChainIntegration as LangChainIntegration + from .litellm import LiteLLMIntegration as LiteLLMIntegration + from .livekit_agents import LiveKitAgentsIntegration as LiveKitAgentsIntegration + from .llamaindex import LlamaIndexIntegration as LlamaIndexIntegration + from .mistral import MistralIntegration as MistralIntegration + from .openai import OpenAIIntegration as OpenAIIntegration + from .openai_agents import OpenAIAgentsIntegration as OpenAIAgentsIntegration + from .openrouter import OpenRouterIntegration as OpenRouterIntegration + from .pipecat import PipecatIntegration as PipecatIntegration + from .pydantic_ai import PydanticAIIntegration as PydanticAIIntegration + from .strands import StrandsIntegration as StrandsIntegration + from .temporal import TemporalIntegration as TemporalIntegration + from .transformers import TransformersIntegration as TransformersIntegration + from .typesafe import TypeSafeIntegration as TypeSafeIntegration __all__ = [ @@ -62,3 +103,16 @@ "TransformersIntegration", "TypeSafeIntegration", ] + + +def __getattr__(name: str) -> Any: + module_name = _INTEGRATION_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f"{__name__}.{module_name}"), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/py/src/braintrust/test_lazy_integration_imports.py b/py/src/braintrust/test_lazy_integration_imports.py new file mode 100644 index 00000000..7672fe30 --- /dev/null +++ b/py/src/braintrust/test_lazy_integration_imports.py @@ -0,0 +1,84 @@ +import json +import subprocess +import sys + + +def _run_import_script(script: str) -> dict[str, object]: + result = subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, text=True) + return json.loads(result.stdout) + + +def test_import_braintrust_does_not_load_integrations_or_providers(): + loaded = _run_import_script( + """ +import importlib.abc +import json +import sys + +class RejectProviderImports(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "openai" or fullname.startswith("openai.") or fullname == "anthropic" or fullname.startswith("anthropic."): + raise AssertionError(f"unexpected provider import: {fullname}") + return None + +sys.meta_path.insert(0, RejectProviderImports()) +import braintrust + +print(json.dumps({ + "integrations": sorted(name for name in sys.modules if name.startswith("braintrust.integrations")), + "providers": sorted(name for name in sys.modules if name == "openai" or name.startswith("openai.") or name == "anthropic" or name.startswith("anthropic.")), +})) +""" + ) + + assert loaded == {"integrations": [], "providers": []} + + +def test_top_level_integration_exports_load_only_when_called(): + loaded = _run_import_script( + """ +import json +import inspect +import sys +import braintrust +from braintrust import setup_ai_sdk, setup_pydantic_ai, wrap_anthropic, wrap_instructor, wrap_litellm, wrap_openai, wrap_openrouter + +before = sorted(name for name in sys.modules if name.startswith("braintrust.integrations")) +braintrust.auto_instrument( + **{name: False for name in inspect.signature(braintrust.auto_instrument).parameters if name != "span_customizers"} +) +after_disabled_auto = sorted(name for name in sys.modules if name.startswith("braintrust.integrations")) + +print(json.dumps({"before": before, "after_disabled_auto": after_disabled_auto})) +""" + ) + + assert loaded == {"before": [], "after_disabled_auto": []} + + +def test_integrations_package_lazily_resolves_public_classes(): + loaded = _run_import_script( + """ +import json +import sys +import braintrust.integrations as integrations + +before = sorted(name for name in sys.modules if name.startswith("braintrust.integrations.")) +integration = integrations.OpenAIIntegration +after = sorted(name for name in sys.modules if name.startswith("braintrust.integrations.")) + +print(json.dumps({ + "before": before, + "resolved_name": integration.__name__, + "loaded_openai": "braintrust.integrations.openai" in after, + "loaded_anthropic": "braintrust.integrations.anthropic" in after, +})) +""" + ) + + assert loaded == { + "before": [], + "resolved_name": "OpenAIIntegration", + "loaded_openai": True, + "loaded_anthropic": False, + } From 90e2c66336880f730f8b5e62f20ea9767891e8bd Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 24 Sep 2026 11:18:03 -0400 Subject: [PATCH 2/2] woops should have updated test --- py/src/braintrust/test_span_customizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py index 0eb5301b..175c6397 100644 --- a/py/src/braintrust/test_span_customizer.py +++ b/py/src/braintrust/test_span_customizer.py @@ -379,7 +379,7 @@ def mask(value): def test_auto_instrument_registration_and_disable(monkeypatch, with_memory_logger, test_logger): # Registration should work without importing optional provider libraries. - monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _: False) + monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class: False) class Redact(SpanCustomizer): def on_span_export(self, data):