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: 1 addition & 1 deletion py/noxfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -255,7 +255,7 @@ def test_dspy(session, version):
session.skip("dspy latest requires Python >= 3.10 (litellm dependency)")
_install_test_deps(session)
_install(session, "dspy", version)
_run_tests(session, f"{WRAPPER_DIR}/test_dspy.py")
_run_tests(session, f"{INTEGRATION_DIR}/dspy/test_dspy.py")


@nox.session()
Expand Down
11 changes: 2 additions & 9 deletions py/src/braintrust/auto.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
AgnoIntegration,
AnthropicIntegration,
ClaudeAgentSDKIntegration,
DSPyIntegration,
GoogleGenAIIntegration,
)

Expand Down Expand Up @@ -125,7 +126,7 @@ def auto_instrument(
if claude_agent_sdk:
results["claude_agent_sdk"] = _instrument_integration(ClaudeAgentSDKIntegration)
if dspy:
results["dspy"] = _instrument_dspy()
results["dspy"] = _instrument_integration(DSPyIntegration)
if adk:
results["adk"] = _instrument_integration(ADKIntegration)

Expand Down Expand Up @@ -160,11 +161,3 @@ def _instrument_pydantic_ai() -> bool:

return setup_pydantic_ai()
return False


def _instrument_dspy() -> bool:
with _try_patch():
from braintrust.wrappers.dspy import patch_dspy

return patch_dspy()
return False
2 changes: 2 additions & 0 deletions py/src/braintrust/integrations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from .agno import AgnoIntegration
from .anthropic import AnthropicIntegration
from .claude_agent_sdk import ClaudeAgentSDKIntegration
from .dspy import DSPyIntegration
from .google_genai import GoogleGenAIIntegration


Expand All @@ -10,5 +11,6 @@
"AgnoIntegration",
"AnthropicIntegration",
"ClaudeAgentSDKIntegration",
"DSPyIntegration",
"GoogleGenAIIntegration",
]
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,15 @@

import dspy
from braintrust.auto import auto_instrument
from braintrust.wrappers.dspy import BraintrustDSpyCallback
from braintrust.integrations.dspy import BraintrustDSpyCallback


# 1. Verify not patched initially
assert not getattr(dspy, "__braintrust_wrapped__", False)
assert not getattr(dspy.configure, "__braintrust_patched_dspy_configure__", False)

# 2. Instrument
results = auto_instrument()
assert results.get("dspy") == True
assert getattr(dspy, "__braintrust_wrapped__", False)

# 3. Idempotent
results2 = auto_instrument()
Expand Down
25 changes: 22 additions & 3 deletions py/src/braintrust/integrations/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,13 +114,27 @@ def patch_marker_attr(cls) -> str:
@classmethod
def mark_patched(cls, obj: Any) -> None:
"""Mark a wrapped target so future patch attempts are idempotent."""
setattr(obj, cls.patch_marker_attr(), True)
try:
setattr(obj, cls.patch_marker_attr(), True)
except AttributeError:
# Some objects (e.g. bound methods) don't support setattr.
# Callers that need a fallback location (like ``patch()``) handle
# this by catching the failure and storing the marker elsewhere.
pass

@classmethod
def is_patched(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> bool:
"""Return whether this patcher's target has already been instrumented."""
marker = cls.patch_marker_attr()
resolved_target = cls.resolve_target(module, version, target=target)
return bool(resolved_target is not None and getattr(resolved_target, cls.patch_marker_attr(), False))
if resolved_target is not None and getattr(resolved_target, marker, False):
return True
# Fall back to checking the root — the marker may live there when the
# resolved target does not support setattr (e.g. bound methods).
root = cls.resolve_root(module, version, target=target)
if root is not None and root is not resolved_target and getattr(root, marker, False):
return True
return False

@classmethod
def patch(cls, module: Any | None, version: str | None, *, target: Any | None = None) -> bool:
Expand All @@ -130,11 +144,16 @@ def patch(cls, module: Any | None, version: str | None, *, target: Any | None =
return False

wrap_function_wrapper(root, cls.target_path, cls.wrapper)
resolved_target = cls.resolve_target(module, version, target=target)
resolved_target = _resolve_attr_path(root, cls.target_path)
if resolved_target is None:
return False

marker = cls.patch_marker_attr()
cls.mark_patched(resolved_target)
# If mark_patched could not store the marker on the target (e.g. bound
# methods), store it on the root so is_patched() can still find it.
if not getattr(resolved_target, marker, False):
setattr(root, marker, True)
return True

@classmethod
Expand Down
12 changes: 12 additions & 0 deletions py/src/braintrust/integrations/dspy/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
"""Braintrust integration for DSPy."""

from .integration import DSPyIntegration
from .patchers import patch_dspy
from .tracing import BraintrustDSpyCallback


__all__ = [
"BraintrustDSpyCallback",
"DSPyIntegration",
"patch_dspy",
]
13 changes: 13 additions & 0 deletions py/src/braintrust/integrations/dspy/integration.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
"""DSPy integration — orchestration class and setup entry-point."""

from braintrust.integrations.base import BaseIntegration

from .patchers import DSPyConfigurePatcher


class DSPyIntegration(BaseIntegration):
"""Braintrust instrumentation for DSPy."""

name = "dspy"
import_names = ("dspy",)
patchers = (DSPyConfigurePatcher,)
43 changes: 43 additions & 0 deletions py/src/braintrust/integrations/dspy/patchers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
"""DSPy patchers — one patcher per coherent patch target."""

from braintrust.integrations.base import FunctionWrapperPatcher

from .tracing import _configure_wrapper


class DSPyConfigurePatcher(FunctionWrapperPatcher):
"""Patch ``dspy.configure`` to auto-add ``BraintrustDSpyCallback``."""

name = "dspy.configure"
target_path = "configure"
wrapper = _configure_wrapper


# ---------------------------------------------------------------------------
# Public helper
# ---------------------------------------------------------------------------


def patch_dspy() -> bool:
"""
Patch DSPy to automatically add Braintrust tracing callback.

After calling this, all calls to dspy.configure() will automatically
include the BraintrustDSpyCallback.

Returns:
True if DSPy was patched (or already patched), False if DSPy is not installed.

Example:
```python
import braintrust
braintrust.patch_dspy()

import dspy
lm = dspy.LM("openai/gpt-4o-mini")
dspy.configure(lm=lm) # BraintrustDSpyCallback auto-added!
```
"""
from .integration import DSPyIntegration

return DSPyIntegration.setup()
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import dspy
import pytest
from braintrust import logger
from braintrust.integrations.dspy import BraintrustDSpyCallback
from braintrust.test_helpers import init_test_logger
from braintrust.wrappers.dspy import BraintrustDSpyCallback
from braintrust.wrappers.test_utils import run_in_subprocess, verify_autoinstrument_script


Expand Down Expand Up @@ -63,17 +63,14 @@ def test_dspy_callback(memory_logger):


class TestPatchDSPy:
"""Tests for patch_dspy() / unpatch_dspy()."""
"""Tests for patch_dspy()."""

def test_patch_dspy_sets_wrapped_flag(self):
"""patch_dspy() should set __braintrust_wrapped__ on dspy module."""
def test_patch_dspy_patches_configure(self):
"""patch_dspy() should patch dspy.configure via the integration patcher."""
result = run_in_subprocess("""
dspy = __import__("dspy")
from braintrust.wrappers.dspy import patch_dspy

assert not hasattr(dspy, "__braintrust_wrapped__")
patch_dspy()
assert hasattr(dspy, "__braintrust_wrapped__")
from braintrust.integrations.dspy import patch_dspy
result = patch_dspy()
assert result, "patch_dspy() should return True"
print("SUCCESS")
""")
assert result.returncode == 0, f"Failed: {result.stderr}"
Expand All @@ -82,7 +79,7 @@ def test_patch_dspy_sets_wrapped_flag(self):
def test_patch_dspy_wraps_configure(self):
"""After patch_dspy(), dspy.configure() should auto-add BraintrustDSpyCallback."""
result = run_in_subprocess("""
from braintrust.wrappers.dspy import patch_dspy, BraintrustDSpyCallback
from braintrust.integrations.dspy import patch_dspy, BraintrustDSpyCallback
patch_dspy()

import dspy
Expand All @@ -103,7 +100,7 @@ def test_patch_dspy_wraps_configure(self):
def test_patch_dspy_preserves_existing_callbacks(self):
"""patch_dspy() should preserve user-provided callbacks."""
result = run_in_subprocess("""
from braintrust.wrappers.dspy import patch_dspy, BraintrustDSpyCallback
from braintrust.integrations.dspy import patch_dspy, BraintrustDSpyCallback
patch_dspy()

import dspy
Expand Down Expand Up @@ -132,7 +129,7 @@ class MyCallback(BaseCallback):
def test_patch_dspy_does_not_duplicate_callback(self):
"""patch_dspy() should not add duplicate BraintrustDSpyCallback."""
result = run_in_subprocess("""
from braintrust.wrappers.dspy import patch_dspy, BraintrustDSpyCallback
from braintrust.integrations.dspy import patch_dspy, BraintrustDSpyCallback
patch_dspy()

import dspy
Expand All @@ -155,7 +152,7 @@ def test_patch_dspy_does_not_duplicate_callback(self):
def test_patch_dspy_idempotent(self):
"""Multiple patch_dspy() calls should be safe."""
result = run_in_subprocess("""
from braintrust.wrappers.dspy import patch_dspy
from braintrust.integrations.dspy import patch_dspy
import dspy

patch_dspy()
Expand All @@ -169,6 +166,17 @@ def test_patch_dspy_idempotent(self):
assert result.returncode == 0, f"Failed: {result.stderr}"
assert "SUCCESS" in result.stdout

def test_legacy_wrapper_import_still_works(self):
"""The old braintrust.wrappers.dspy import path should still work."""
result = run_in_subprocess("""
from braintrust.wrappers.dspy import BraintrustDSpyCallback, patch_dspy
assert BraintrustDSpyCallback is not None
assert callable(patch_dspy)
print("SUCCESS")
""")
assert result.returncode == 0, f"Failed: {result.stderr}"
assert "SUCCESS" in result.stdout


class TestAutoInstrumentDSPy:
"""Tests for auto_instrument() with DSPy."""
Expand Down
Loading
Loading