fix: Lazy load integration modules

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)

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

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__))