From cbf491d3ea918da582fc49cd53c4b50e8e1de9ea Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 31 Mar 2026 17:02:39 -0400 Subject: [PATCH] feat(langchain): add callback patcher integration Add a CallbackPatcher base primitive and use it to model LangChain as an SDK integration that installs a default Braintrust callback handler. Export the LangChain integration, wire it into auto_instrument(), and add focused setup and auto-instrument tests. Also document CallbackPatcher usage in the sdk-integrations skill so future integrations can use the same pattern for setup-only side effects. --- .agents/skills/sdk-integrations/SKILL.md | 2 + py/src/braintrust/auto.py | 5 ++ py/src/braintrust/integrations/__init__.py | 2 + .../auto_test_scripts/test_auto_langchain.py | 55 ++++++++++++++++ py/src/braintrust/integrations/base.py | 63 +++++++++++++++++++ .../integrations/langchain/__init__.py | 40 +++++++++++- .../integrations/langchain/callbacks.py | 36 +++++++---- .../integrations/langchain/context.py | 6 +- .../integrations/langchain/integration.py | 13 ++++ .../integrations/langchain/patchers.py | 37 +++++++++++ .../integrations/langchain/test_context.py | 26 +++++++- 11 files changed, 267 insertions(+), 18 deletions(-) create mode 100644 py/src/braintrust/integrations/auto_test_scripts/test_auto_langchain.py create mode 100644 py/src/braintrust/integrations/langchain/integration.py create mode 100644 py/src/braintrust/integrations/langchain/patchers.py diff --git a/.agents/skills/sdk-integrations/SKILL.md b/.agents/skills/sdk-integrations/SKILL.md index 7a9ca12e..92eed335 100644 --- a/.agents/skills/sdk-integrations/SKILL.md +++ b/.agents/skills/sdk-integrations/SKILL.md @@ -224,6 +224,8 @@ Use `CompositeFunctionWrapperPatcher` when several closely related targets shoul - sync and async variants of the same method - the same logical surface patched across multiple modules +Use `CallbackPatcher` when the integration only needs a setup side effect after applicability succeeds, not in-place wrapping or class replacement. Good fits include registering a global callback handler, invoking a provider-owned configure hook, or installing default integration state. Gate it with `target_module` when the side effect depends on an optional module, and use `state_getter` when idempotence should be derived from integration-managed state instead of a marker on the imported root. + Set `target_module` when the patch target lives outside the module named by `import_names`, especially for optional or deep submodules. Failed `target_module` imports should make the patcher skip cleanly through `applies()`. Use `superseded_by` for version-conditional mutual exclusion. Express fallback relationships declaratively instead of reproducing `hasattr` logic in custom `applies()` methods whenever possible. diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index 55e07e84..dc44c7d2 100644 --- a/py/src/braintrust/auto.py +++ b/py/src/braintrust/auto.py @@ -15,6 +15,7 @@ ClaudeAgentSDKIntegration, DSPyIntegration, GoogleGenAIIntegration, + LangChainIntegration, LiteLLMIntegration, OpenRouterIntegration, PydanticAIIntegration, @@ -50,6 +51,7 @@ def auto_instrument( claude_agent_sdk: bool = True, dspy: bool = True, adk: bool = True, + langchain: bool = True, ) -> dict[str, bool]: """ Auto-instrument supported AI/ML libraries for Braintrust tracing. @@ -72,6 +74,7 @@ def auto_instrument( claude_agent_sdk: Enable Claude Agent SDK instrumentation (default: True) dspy: Enable DSPy instrumentation (default: True) adk: Enable Google ADK instrumentation (default: True) + langchain: Enable LangChain instrumentation (default: True) Returns: Dict mapping integration name to whether it was successfully instrumented. @@ -141,6 +144,8 @@ def auto_instrument( results["dspy"] = _instrument_integration(DSPyIntegration) if adk: results["adk"] = _instrument_integration(ADKIntegration) + if langchain: + results["langchain"] = _instrument_integration(LangChainIntegration) return results diff --git a/py/src/braintrust/integrations/__init__.py b/py/src/braintrust/integrations/__init__.py index 562f9ae1..0062ec77 100644 --- a/py/src/braintrust/integrations/__init__.py +++ b/py/src/braintrust/integrations/__init__.py @@ -5,6 +5,7 @@ from .claude_agent_sdk import ClaudeAgentSDKIntegration from .dspy import DSPyIntegration from .google_genai import GoogleGenAIIntegration +from .langchain import LangChainIntegration from .litellm import LiteLLMIntegration from .openrouter import OpenRouterIntegration from .pydantic_ai import PydanticAIIntegration @@ -19,6 +20,7 @@ "DSPyIntegration", "GoogleGenAIIntegration", "LiteLLMIntegration", + "LangChainIntegration", "OpenRouterIntegration", "PydanticAIIntegration", ] diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_langchain.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_langchain.py new file mode 100644 index 00000000..e7b6371f --- /dev/null +++ b/py/src/braintrust/integrations/auto_test_scripts/test_auto_langchain.py @@ -0,0 +1,55 @@ +"""Test auto_instrument for LangChain.""" + +from pathlib import Path + +from braintrust.auto import auto_instrument +from braintrust.integrations.langchain import BraintrustCallbackHandler +from braintrust.integrations.langchain.context import clear_global_handler, get_global_handler +from braintrust.wrappers.test_utils import autoinstrument_test_context +from langchain_core.callbacks import CallbackManager +from langchain_core.prompts import ChatPromptTemplate +from langchain_openai import ChatOpenAI + + +_CASSETTES_DIR = Path(__file__).resolve().parent.parent / "langchain" / "cassettes" + +# 1. Verify not patched initially. +clear_global_handler() +assert get_global_handler() is None +manager = CallbackManager.configure() +assert next((h for h in manager.handlers if isinstance(h, BraintrustCallbackHandler)), None) is None + +# 2. Instrument. +results = auto_instrument() +assert results.get("langchain") == True +handler = get_global_handler() +assert isinstance(handler, BraintrustCallbackHandler) + +manager = CallbackManager.configure() +assert next((h for h in manager.handlers if isinstance(h, BraintrustCallbackHandler)), None) is handler + +# 3. Idempotent. +results2 = auto_instrument() +assert results2.get("langchain") == True +assert get_global_handler() is handler + +# 4. Make an API call and verify spans. +with autoinstrument_test_context("test_global_handler", cassettes_dir=_CASSETTES_DIR) as memory_logger: + prompt = ChatPromptTemplate.from_template("What is 1 + {number}?") + model = ChatOpenAI( + model="gpt-4o-mini", + temperature=1, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + n=1, + ) + chain = prompt.pipe(model) + + message = chain.invoke({"number": "2"}) + assert message.content == "1 + 2 equals 3." + + spans = memory_logger.pop() + assert len(spans) > 0 + +print("SUCCESS") diff --git a/py/src/braintrust/integrations/base.py b/py/src/braintrust/integrations/base.py index 62e0c8ce..68d2936a 100644 --- a/py/src/braintrust/integrations/base.py +++ b/py/src/braintrust/integrations/base.py @@ -69,6 +69,69 @@ def patch(cls, module: Any | None, version: str | None, *, target: Any | None = raise NotImplementedError +class CallbackPatcher(BasePatcher): + """Base patcher for integration setup steps that only execute a callback. + + Use this for integrations that do not need in-place wrapping or class + replacement, but still need an idempotent setup action once a target module + is available — for example, registering a global callback handler. + + Set ``target_module`` when the callback should only run if a particular + optional module can be imported. Provide ``state_getter`` when patch state + should be derived from integration-owned state instead of a marker stored on + the resolved root object. + """ + + target_module: ClassVar[str | None] = None + callback: ClassVar[Any] + state_getter: ClassVar[Any | None] = None + + @classmethod + def resolve_root(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> Any | None: + """Return the object whose availability gates this callback patcher.""" + if target is not None: + return target + if cls.target_module is not None: + try: + return importlib.import_module(cls.target_module) + except ImportError: + return None + return module + + @classmethod + def applies(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> bool: + """Return whether the callback should run for the given module/version.""" + return ( + super().applies(module, version, target=target) + and cls.resolve_root(module, version, target=target) is not None + ) + + @classmethod + def is_patched(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> bool: + """Return whether this callback patcher has already been applied.""" + if cls.state_getter is not None: + return bool(cls.state_getter()) + root = cls.resolve_root(module, version, target=target) + return bool(root is not None and cls.has_patch_marker(root)) + + @classmethod + def patch(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> bool: + """Execute the callback and mark the root as patched when needed.""" + root = cls.resolve_root(module, version, target=target) + if root is None or not cls.applies(module, version, target=target): + return False + + result = cls.callback() + if result is False: + return False + + if cls.state_getter is not None: + return cls.is_patched(module, version, target=target) + + cls.mark_patched(root) + return True + + class ClassScanPatcher(BasePatcher): """Base patcher for rescanning and patching discovered class hierarchies.""" diff --git a/py/src/braintrust/integrations/langchain/__init__.py b/py/src/braintrust/integrations/langchain/__init__.py index 67ecb04a..f2a0aee7 100644 --- a/py/src/braintrust/integrations/langchain/__init__.py +++ b/py/src/braintrust/integrations/langchain/__init__.py @@ -1,5 +1,39 @@ -from braintrust.integrations.langchain.callbacks import BraintrustCallbackHandler -from braintrust.integrations.langchain.context import set_global_handler +"""Braintrust integration for LangChain.""" +from braintrust.logger import NOOP_SPAN, current_span, init_logger -__all__ = ["BraintrustCallbackHandler", "set_global_handler"] +from .integration import LangChainIntegration + + +try: + from .callbacks import BraintrustCallbackHandler + from .context import set_global_handler +except ImportError as exc: # pragma: no cover - optional dependency not installed + _IMPORT_ERROR = exc + + class BraintrustCallbackHandler: # type: ignore[no-redef] + def __init__(self, *args, **kwargs): + raise ImportError("langchain-core is required for braintrust.integrations.langchain") from _IMPORT_ERROR + + def set_global_handler(handler): # type: ignore[no-redef] + raise ImportError("langchain-core is required for braintrust.integrations.langchain") from _IMPORT_ERROR + + +__all__ = [ + "BraintrustCallbackHandler", + "LangChainIntegration", + "set_global_handler", + "setup_langchain", +] + + +def setup_langchain( + api_key: str | None = None, + project_id: str | None = None, + project_name: str | None = None, +) -> bool: + """Setup Braintrust integration with LangChain.""" + if current_span() == NOOP_SPAN: + init_logger(project=project_name, api_key=api_key, project_id=project_id) + + return LangChainIntegration.setup() diff --git a/py/src/braintrust/integrations/langchain/callbacks.py b/py/src/braintrust/integrations/langchain/callbacks.py index e54bb365..50650da0 100644 --- a/py/src/braintrust/integrations/langchain/callbacks.py +++ b/py/src/braintrust/integrations/langchain/callbacks.py @@ -11,8 +11,9 @@ ) from uuid import UUID -import braintrust -from braintrust import NOOP_SPAN, Logger, Span, SpanAttributes, SpanTypeAttribute, current_span, init_logger +from braintrust.generated_types import SpanAttributes +from braintrust.logger import NOOP_SPAN, Logger, Span, current_span, init_logger, start_span +from braintrust.span_types import SpanTypeAttribute from braintrust.version import VERSION as sdk_version from langchain_core.agents import AgentAction, AgentFinish from langchain_core.callbacks.base import BaseCallbackHandler @@ -93,8 +94,6 @@ def _start_span( parent_span = current_parent elif self.logger is not None: parent_span = self.logger - else: - parent_span = braintrust if event is None: event = {} @@ -116,15 +115,26 @@ def _start_span( }, } - span = parent_span.start_span( - name=name, - type=type, - span_attributes=span_attributes, - start_time=start_time, - set_current=set_current, - parent=parent, - **event, - ) + if parent_span is None: + span = start_span( + name=name, + type=type, + span_attributes=span_attributes, + start_time=start_time, + set_current=set_current, + parent=parent, + **event, + ) + else: + span = parent_span.start_span( + name=name, + type=type, + span_attributes=span_attributes, + start_time=start_time, + set_current=set_current, + parent=parent, + **event, + ) if self.logger != NOOP_SPAN and span == NOOP_SPAN: _logger.warning( diff --git a/py/src/braintrust/integrations/langchain/context.py b/py/src/braintrust/integrations/langchain/context.py index 9de61670..120e0871 100644 --- a/py/src/braintrust/integrations/langchain/context.py +++ b/py/src/braintrust/integrations/langchain/context.py @@ -4,7 +4,7 @@ from langchain_core.tracers.context import register_configure_hook -__all__ = ["set_global_handler", "clear_global_handler"] +__all__ = ["set_global_handler", "clear_global_handler", "get_global_handler"] braintrust_callback_handler_var: ContextVar[BraintrustCallbackHandler | None] = ContextVar( @@ -16,6 +16,10 @@ def set_global_handler(handler: BraintrustCallbackHandler): braintrust_callback_handler_var.set(handler) +def get_global_handler() -> BraintrustCallbackHandler | None: + return braintrust_callback_handler_var.get() + + def clear_global_handler(): braintrust_callback_handler_var.set(None) diff --git a/py/src/braintrust/integrations/langchain/integration.py b/py/src/braintrust/integrations/langchain/integration.py new file mode 100644 index 00000000..877c28c3 --- /dev/null +++ b/py/src/braintrust/integrations/langchain/integration.py @@ -0,0 +1,13 @@ +"""LangChain integration definition.""" + +from braintrust.integrations.base import BaseIntegration + +from .patchers import GlobalHandlerPatcher + + +class LangChainIntegration(BaseIntegration): + """Braintrust instrumentation for LangChain via a global callback handler.""" + + name = "langchain" + import_names = ("langchain_core",) + patchers = (GlobalHandlerPatcher,) diff --git a/py/src/braintrust/integrations/langchain/patchers.py b/py/src/braintrust/integrations/langchain/patchers.py new file mode 100644 index 00000000..1e408bc2 --- /dev/null +++ b/py/src/braintrust/integrations/langchain/patchers.py @@ -0,0 +1,37 @@ +"""LangChain patchers.""" + +from braintrust.integrations.base import CallbackPatcher + + +try: + from .callbacks import BraintrustCallbackHandler + from .context import get_global_handler, set_global_handler +except ImportError: # pragma: no cover - optional dependency not installed + BraintrustCallbackHandler = None + get_global_handler = None + set_global_handler = None + + +def _has_global_handler() -> bool: + if BraintrustCallbackHandler is None or get_global_handler is None: + return False + return isinstance(get_global_handler(), BraintrustCallbackHandler) + + +def _set_default_global_handler() -> None: + if BraintrustCallbackHandler is None or get_global_handler is None or set_global_handler is None: + raise ImportError("langchain_core is not installed") + + if isinstance(get_global_handler(), BraintrustCallbackHandler): + return + + set_global_handler(BraintrustCallbackHandler()) + + +class GlobalHandlerPatcher(CallbackPatcher): + """Install the default Braintrust callback handler into LangChain's configure hook.""" + + name = "langchain.global_handler" + target_module = "langchain_core.tracers.context" + callback = _set_default_global_handler + state_getter = _has_global_handler diff --git a/py/src/braintrust/integrations/langchain/test_context.py b/py/src/braintrust/integrations/langchain/test_context.py index 6d29a959..a047ad41 100644 --- a/py/src/braintrust/integrations/langchain/test_context.py +++ b/py/src/braintrust/integrations/langchain/test_context.py @@ -4,8 +4,9 @@ import pytest from braintrust import logger -from braintrust.integrations.langchain import BraintrustCallbackHandler, set_global_handler +from braintrust.integrations.langchain import BraintrustCallbackHandler, set_global_handler, setup_langchain from braintrust.test_helpers import init_test_logger +from braintrust.wrappers.test_utils import verify_autoinstrument_script from langchain_core.callbacks import CallbackManager from langchain_core.messages import BaseMessage from langchain_core.prompts import ChatPromptTemplate @@ -153,3 +154,26 @@ def test_global_handler(logger_memory_logger): ) assert message.content == "1 + 2 equals 3." + + +def test_setup_langchain_installs_default_handler(): + from braintrust.integrations.langchain.context import get_global_handler + + manager = CallbackManager.configure() + assert next((h for h in manager.handlers if isinstance(h, BraintrustCallbackHandler)), None) is None + assert get_global_handler() is None + + assert setup_langchain() + handler = get_global_handler() + assert isinstance(handler, BraintrustCallbackHandler) + + manager = CallbackManager.configure() + assert next((h for h in manager.handlers if isinstance(h, BraintrustCallbackHandler)), None) is handler + + assert setup_langchain() + assert get_global_handler() is handler + + +class TestAutoInstrumentLangChain: + def test_auto_instrument_langchain(self): + verify_autoinstrument_script("test_auto_langchain.py")