Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .agents/skills/sdk-integrations/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions py/src/braintrust/auto.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
ClaudeAgentSDKIntegration,
DSPyIntegration,
GoogleGenAIIntegration,
LangChainIntegration,
LiteLLMIntegration,
OpenRouterIntegration,
PydanticAIIntegration,
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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

Expand Down
2 changes: 2 additions & 0 deletions py/src/braintrust/integrations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -19,6 +20,7 @@
"DSPyIntegration",
"GoogleGenAIIntegration",
"LiteLLMIntegration",
"LangChainIntegration",
"OpenRouterIntegration",
"PydanticAIIntegration",
]
Original file line number Diff line number Diff line change
@@ -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")
63 changes: 63 additions & 0 deletions py/src/braintrust/integrations/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
40 changes: 37 additions & 3 deletions py/src/braintrust/integrations/langchain/__init__.py
Original file line number Diff line number Diff line change
@@ -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()
36 changes: 23 additions & 13 deletions py/src/braintrust/integrations/langchain/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 = {}
Expand All @@ -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(
Expand Down
6 changes: 5 additions & 1 deletion py/src/braintrust/integrations/langchain/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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)

Expand Down
13 changes: 13 additions & 0 deletions py/src/braintrust/integrations/langchain/integration.py
Original file line number Diff line number Diff line change
@@ -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,)
37 changes: 37 additions & 0 deletions py/src/braintrust/integrations/langchain/patchers.py
Original file line number Diff line number Diff line change
@@ -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
26 changes: 25 additions & 1 deletion py/src/braintrust/integrations/langchain/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Loading