braintrust 0.4.3__py3-none-any.whl → 0.5.2__py3-none-any.whl

braintrust/wrappers/claude_agent_sdk/test_wrapper.py

@@ -23,6 +23,7 @@ from braintrust.wrappers.claude_agent_sdk._wrapper import (
     _create_client_wrapper_class,
     _create_tool_wrapper_class,
 )
+from braintrust.wrappers.test_utils import verify_autoinstrument_script
 
 PROJECT_NAME = "test-claude-agent-sdk"
 TEST_MODEL = "claude-haiku-4-5-20251001"
@@ -177,3 +178,117 @@ async def test_calculator_with_multiple_operations(memory_logger):
         if span["span_id"] != root_span_id:
             assert span["root_span_id"] == root_span_id
             assert root_span_id in span["span_parents"]
+
+
+def _make_message(content: str) -> dict:
+    """Create a streaming format message dict."""
+    return {"type": "user", "message": {"role": "user", "content": content}}
+
+
+def _assert_structured_input(task_span: dict, expected_contents: list[str]) -> None:
+    """Assert that task span input is a structured list with expected content."""
+    inp = task_span.get("input")
+    assert isinstance(inp, list), f"Expected list input, got {type(inp).__name__}: {inp}"
+    assert [x["message"]["content"] for x in inp] == expected_contents
+
+
+class CustomAsyncIterable:
+    """Custom AsyncIterable class (not a generator) for testing."""
+
+    def __init__(self, messages: list[dict]):
+        self._messages = messages
+
+    def __aiter__(self):
+        return CustomAsyncIterator(self._messages)
+
+
+class CustomAsyncIterator:
+    """Iterator for CustomAsyncIterable."""
+
+    def __init__(self, messages: list[dict]):
+        self._messages = messages
+        self._index = 0
+
+    async def __anext__(self):
+        if self._index >= len(self._messages):
+            raise StopAsyncIteration
+        msg = self._messages[self._index]
+        self._index += 1
+        return msg
+
+
+@pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed")
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "input_factory,expected_contents",
+    [
+        pytest.param(
+            lambda: (msg async for msg in _single_message_generator()),
+            ["What is 2 + 2?"],
+            id="asyncgen_single",
+        ),
+        pytest.param(
+            lambda: (msg async for msg in _multi_message_generator()),
+            ["Part 1", "Part 2"],
+            id="asyncgen_multi",
+        ),
+        pytest.param(
+            lambda: CustomAsyncIterable([_make_message("Custom 1"), _make_message("Custom 2")]),
+            ["Custom 1", "Custom 2"],
+            id="custom_async_iterable",
+        ),
+    ],
+)
+async def test_query_async_iterable(memory_logger, input_factory, expected_contents):
+    """Test that async iterable inputs are captured as structured lists.
+
+    Verifies that passing AsyncIterable[dict] to query() results in the span
+    input showing the structured message list, not a flattened string or repr.
+    """
+    assert not memory_logger.pop()
+
+    original_client = claude_agent_sdk.ClaudeSDKClient
+    claude_agent_sdk.ClaudeSDKClient = _create_client_wrapper_class(original_client)
+
+    try:
+        options = claude_agent_sdk.ClaudeAgentOptions(model=TEST_MODEL)
+
+        async with claude_agent_sdk.ClaudeSDKClient(options=options) as client:
+            await client.query(input_factory())
+            async for message in client.receive_response():
+                if type(message).__name__ == "ResultMessage":
+                    break
+
+        spans = memory_logger.pop()
+
+        task_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.TASK]
+        assert len(task_spans) >= 1, f"Should have at least one task span, got {len(task_spans)}"
+
+        task_span = next(
+            (s for s in task_spans if s["span_attributes"]["name"] == "Claude Agent"),
+            task_spans[0],
+        )
+
+        _assert_structured_input(task_span, expected_contents)
+
+    finally:
+        claude_agent_sdk.ClaudeSDKClient = original_client
+
+
+async def _single_message_generator():
+    """Generator yielding a single message."""
+    yield _make_message("What is 2 + 2?")
+
+
+async def _multi_message_generator():
+    """Generator yielding multiple messages."""
+    yield _make_message("Part 1")
+    yield _make_message("Part 2")
+
+
+class TestAutoInstrumentClaudeAgentSDK:
+    """Tests for auto_instrument() with Claude Agent SDK."""
+
+    def test_auto_instrument_claude_agent_sdk(self):
+        """Test auto_instrument patches Claude Agent SDK and creates spans."""
+        verify_autoinstrument_script("test_auto_claude_agent_sdk.py")

braintrust/wrappers/dspy.py

@@ -51,6 +51,7 @@ from typing import Any
 
 from braintrust.logger import current_span, start_span
 from braintrust.span_types import SpanTypeAttribute
+from wrapt import wrap_function_wrapper
 
 # Note: For detailed token and cost metrics, use patch_litellm() before importing DSPy.
 # The DSPy callback focuses on execution flow and span hierarchy.
@@ -60,6 +61,8 @@ try:
 except ImportError:
     raise ImportError("DSPy is not installed. Please install it with: pip install dspy")
 
+__all__ = ["BraintrustDSpyCallback", "patch_dspy"]
+
 
 class BraintrustDSpyCallback(BaseCallback):
     """Callback handler that logs DSPy execution traces to Braintrust.
@@ -412,4 +415,52 @@ class BraintrustDSpyCallback(BaseCallback):
             span.end()
 
 
-__all__ = ["BraintrustDSpyCallback"]
+def _configure_wrapper(wrapped, instance, args, kwargs):
+    """Wrapper for dspy.configure that auto-adds BraintrustDSpyCallback."""
+    callbacks = kwargs.get("callbacks")
+    if callbacks is None:
+        callbacks = []
+    else:
+        callbacks = list(callbacks)
+
+    # Check if already has Braintrust callback
+    has_bt_callback = any(isinstance(cb, BraintrustDSpyCallback) for cb in callbacks)
+    if not has_bt_callback:
+        callbacks.append(BraintrustDSpyCallback())
+
+    kwargs["callbacks"] = callbacks
+    return wrapped(*args, **kwargs)
+
+
+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!
+        ```
+    """
+    try:
+        import dspy
+
+        if getattr(dspy, "__braintrust_wrapped__", False):
+            return True  # Already patched
+
+        wrap_function_wrapper("dspy", "configure", _configure_wrapper)
+        dspy.__braintrust_wrapped__ = True
+        return True
+
+    except ImportError:
+        return False

braintrust/wrappers/google_genai/__init__.py

@@ -15,7 +15,13 @@ def setup_genai(
     api_key: str | None = None,
     project_id: str | None = None,
     project_name: str | None = None,
-):
+) -> bool:
+    """
+    Setup Braintrust integration with Google GenAI.
+
+    Returns:
+        True if setup was successful, False if google-genai is not installed.
+    """
     span = current_span()
     if span == NOOP_SPAN:
         init_logger(project=project_name, api_key=api_key, project_id=project_id)
@@ -27,11 +33,8 @@ def setup_genai(
         genai.Client = wrap_client(genai.Client)
         models.Models = wrap_models(models.Models)
         models.AsyncModels = wrap_async_models(models.AsyncModels)
-        pass
-    except ImportError as e:
-        logger.error(
-            f"Failed to import Google ADK agents: {e}. Google ADK is not installed. Please install it with: pip install google-adk"
-        )
+        return True
+    except ImportError:
         return False
 
 

braintrust/wrappers/litellm.py

@@ -631,13 +631,16 @@ def serialize_response_format(response_format: Any) -> Any:
         return response_format
 
 
-def patch_litellm():
+def patch_litellm() -> bool:
     """
     Patch LiteLLM to add Braintrust tracing.
 
     This wraps litellm.completion and litellm.acompletion to automatically
     create Braintrust spans with detailed token metrics, timing, and costs.
 
+    Returns:
+        True if LiteLLM was patched (or already patched), False if LiteLLM is not installed.
+
     Example:
         ```python
         import braintrust
@@ -657,52 +660,12 @@ def patch_litellm():
         import litellm
 
         if not hasattr(litellm, "_braintrust_wrapped"):
-            # Store originals for unpatch_litellm()
-            litellm._braintrust_original_completion = litellm.completion
-            litellm._braintrust_original_acompletion = litellm.acompletion
-            litellm._braintrust_original_responses = litellm.responses
-            litellm._braintrust_original_aresponses = litellm.aresponses
-
             wrapped = wrap_litellm(litellm)
             litellm.completion = wrapped.completion
             litellm.acompletion = wrapped.acompletion
             litellm.responses = wrapped.responses
             litellm.aresponses = wrapped.aresponses
             litellm._braintrust_wrapped = True
+        return True
     except ImportError:
-        pass  # litellm not available
-
-
-def unpatch_litellm():
-    """
-    Restore LiteLLM to its original state, removing Braintrust tracing.
-
-    This undoes the patching done by patch_litellm(), restoring the original
-    completion, acompletion, responses, and aresponses functions.
-
-    Example:
-        ```python
-        import braintrust
-        braintrust.patch_litellm()
-
-        # ... use litellm with tracing ...
-
-        braintrust.unpatch_litellm()  # restore original behavior
-        ```
-    """
-    try:
-        import litellm
-
-        if hasattr(litellm, "_braintrust_wrapped"):
-            litellm.completion = litellm._braintrust_original_completion
-            litellm.acompletion = litellm._braintrust_original_acompletion
-            litellm.responses = litellm._braintrust_original_responses
-            litellm.aresponses = litellm._braintrust_original_aresponses
-
-            delattr(litellm, "_braintrust_wrapped")
-            delattr(litellm, "_braintrust_original_completion")
-            delattr(litellm, "_braintrust_original_acompletion")
-            delattr(litellm, "_braintrust_original_responses")
-            delattr(litellm, "_braintrust_original_aresponses")
-    except ImportError:
-        pass  # litellm not available
+        return False

braintrust/wrappers/pydantic_ai.py

@@ -51,9 +51,8 @@ def setup_pydantic_ai(
         wrap_model_classes()
 
         return True
-    except ImportError as e:
-        logger.error(f"Failed to import Pydantic AI: {e}")
-        logger.error("Pydantic AI is not installed. Please install it with: pip install pydantic-ai-slim")
+    except ImportError:
+        # Not installed - this is expected when using auto_instrument()
         return False
 
 

braintrust/wrappers/test_agno.py

@@ -8,6 +8,7 @@ import pytest
 from braintrust import logger
 from braintrust.test_helpers import init_test_logger
 from braintrust.wrappers.agno import setup_agno
+from braintrust.wrappers.test_utils import verify_autoinstrument_script
 
 TEST_ORG_ID = "test-org-123"
 PROJECT_NAME = "test-agno-app"
@@ -94,3 +95,11 @@ def test_agno_simple_agent_execution(memory_logger):
     assert llm_span["metrics"]["prompt_tokens"] == 38
     assert llm_span["metrics"]["completion_tokens"] == 4
     assert llm_span["metrics"]["tokens"] == 42
+
+
+class TestAutoInstrumentAgno:
+    """Tests for auto_instrument() with Agno."""
+
+    def test_auto_instrument_agno(self):
+        """Test auto_instrument patches Agno and creates spans."""
+        verify_autoinstrument_script("test_auto_agno.py")

braintrust/wrappers/test_anthropic.py

@@ -9,6 +9,7 @@ import pytest
 from braintrust import logger
 from braintrust.test_helpers import init_test_logger
 from braintrust.wrappers.anthropic import wrap_anthropic
+from braintrust.wrappers.test_utils import run_in_subprocess, verify_autoinstrument_script
 
 TEST_ORG_ID = "test-org-123"
 PROJECT_NAME = "test-anthropic-app"
@@ -481,3 +482,158 @@ async def test_anthropic_beta_messages_streaming_async(memory_logger):
     assert metrics["prompt_tokens"] == usage.input_tokens
     assert metrics["completion_tokens"] == usage.output_tokens
     assert metrics["tokens"] == usage.input_tokens + usage.output_tokens
+
+
+class TestPatchAnthropic:
+    """Tests for patch_anthropic() / unpatch_anthropic()."""
+
+    def test_patch_anthropic_sets_wrapped_flag(self):
+        """patch_anthropic() should set __braintrust_wrapped__ on anthropic module."""
+        result = run_in_subprocess("""
+            from braintrust.wrappers.anthropic import patch_anthropic
+            import anthropic
+
+            assert not hasattr(anthropic, "__braintrust_wrapped__")
+            patch_anthropic()
+            assert hasattr(anthropic, "__braintrust_wrapped__")
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    def test_patch_anthropic_wraps_new_clients(self):
+        """After patch_anthropic(), new Anthropic() clients should be wrapped."""
+        result = run_in_subprocess("""
+            from braintrust.wrappers.anthropic import patch_anthropic
+            patch_anthropic()
+
+            import anthropic
+            client = anthropic.Anthropic(api_key="test-key")
+
+            # Check that messages is wrapped
+            messages_type = type(client.messages).__name__
+            print(f"messages_type={messages_type}")
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    def test_patch_anthropic_idempotent(self):
+        """Multiple patch_anthropic() calls should be safe."""
+        result = run_in_subprocess("""
+            from braintrust.wrappers.anthropic import patch_anthropic
+            import anthropic
+
+            patch_anthropic()
+            first_class = anthropic.Anthropic
+
+            patch_anthropic()  # Second call
+            second_class = anthropic.Anthropic
+
+            assert first_class is second_class
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    def test_patch_anthropic_creates_spans(self):
+        """patch_anthropic() should create spans when making API calls."""
+        result = run_in_subprocess("""
+            from braintrust.wrappers.anthropic import patch_anthropic
+            from braintrust.test_helpers import init_test_logger
+            from braintrust import logger
+
+            # Set up memory logger
+            init_test_logger("test-auto")
+            with logger._internal_with_memory_background_logger() as memory_logger:
+                patch_anthropic()
+
+                import anthropic
+                client = anthropic.Anthropic()
+
+                # Make a call within a span context
+                import braintrust
+                with braintrust.start_span(name="test") as span:
+                    try:
+                        # This will fail without API key, but span should still be created
+                        client.messages.create(
+                            model="claude-3-5-haiku-latest",
+                            max_tokens=100,
+                            messages=[{"role": "user", "content": "hi"}],
+                        )
+                    except Exception:
+                        pass  # Expected without API key
+
+                # Check that spans were logged
+                spans = memory_logger.pop()
+                # Should have at least the parent span
+                assert len(spans) >= 1, f"Expected spans, got {spans}"
+                print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+
+class TestPatchAnthropicSpans:
+    """VCR-based tests verifying that patch_anthropic() produces spans."""
+
+    @pytest.mark.vcr
+    def test_patch_anthropic_creates_spans(self, memory_logger):
+        """patch_anthropic() should create spans when making API calls."""
+        from braintrust.wrappers.anthropic import patch_anthropic
+
+        assert not memory_logger.pop()
+
+        patch_anthropic()
+        client = anthropic.Anthropic()
+        response = client.messages.create(
+            model="claude-3-5-haiku-latest",
+            max_tokens=100,
+            messages=[{"role": "user", "content": "Say hi"}],
+        )
+        assert response.content[0].text
+
+        # Verify span was created
+        spans = memory_logger.pop()
+        assert len(spans) == 1
+        span = spans[0]
+        assert span["metadata"]["provider"] == "anthropic"
+        assert "claude" in span["metadata"]["model"]
+        assert span["input"]
+
+
+class TestPatchAnthropicAsyncSpans:
+    """VCR-based tests verifying that patch_anthropic() produces spans for async clients."""
+
+    @pytest.mark.vcr
+    @pytest.mark.asyncio
+    async def test_patch_anthropic_async_creates_spans(self, memory_logger):
+        """patch_anthropic() should create spans for async API calls."""
+        from braintrust.wrappers.anthropic import patch_anthropic
+
+        assert not memory_logger.pop()
+
+        patch_anthropic()
+        client = anthropic.AsyncAnthropic()
+        response = await client.messages.create(
+            model="claude-3-5-haiku-latest",
+            max_tokens=100,
+            messages=[{"role": "user", "content": "Say hi async"}],
+        )
+        assert response.content[0].text
+
+        # Verify span was created
+        spans = memory_logger.pop()
+        assert len(spans) == 1
+        span = spans[0]
+        assert span["metadata"]["provider"] == "anthropic"
+        assert "claude" in span["metadata"]["model"]
+        assert span["input"]
+
+
+class TestAutoInstrumentAnthropic:
+    """Tests for auto_instrument() with Anthropic."""
+
+    def test_auto_instrument_anthropic(self):
+        """Test auto_instrument patches Anthropic, creates spans, and uninstrument works."""
+        verify_autoinstrument_script("test_auto_anthropic.py")

braintrust/wrappers/test_dspy.py

@@ -7,6 +7,7 @@ import pytest
 from braintrust import logger
 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
 
 PROJECT_NAME = "test-dspy-app"
 MODEL = "openai/gpt-4o-mini"
@@ -58,3 +59,119 @@ def test_dspy_callback(memory_logger):
 
     # Verify span parenting (LM span should have parent)
     assert lm_span.get("span_parents")  # LM span should have parent
+
+
+class TestPatchDSPy:
+    """Tests for patch_dspy() / unpatch_dspy()."""
+
+    def test_patch_dspy_sets_wrapped_flag(self):
+        """patch_dspy() should set __braintrust_wrapped__ on dspy module."""
+        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__")
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    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
+            patch_dspy()
+
+            import dspy
+
+            # Configure without explicitly adding callback
+            dspy.configure(lm=None)
+
+            # Check that BraintrustDSpyCallback was auto-added
+            from dspy.dsp.utils.settings import settings
+            callbacks = settings.callbacks
+            has_bt_callback = any(isinstance(cb, BraintrustDSpyCallback) for cb in callbacks)
+            assert has_bt_callback, f"Expected BraintrustDSpyCallback in {callbacks}"
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    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
+            patch_dspy()
+
+            import dspy
+            from dspy.utils.callback import BaseCallback
+
+            class MyCallback(BaseCallback):
+                pass
+
+            my_callback = MyCallback()
+            dspy.configure(lm=None, callbacks=[my_callback])
+
+            from dspy.dsp.utils.settings import settings
+            callbacks = settings.callbacks
+
+            # Should have both callbacks
+            has_my_callback = any(cb is my_callback for cb in callbacks)
+            has_bt_callback = any(isinstance(cb, BraintrustDSpyCallback) for cb in callbacks)
+
+            assert has_my_callback, "User callback should be preserved"
+            assert has_bt_callback, "BraintrustDSpyCallback should be added"
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    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
+            patch_dspy()
+
+            import dspy
+
+            # User explicitly adds BraintrustDSpyCallback
+            bt_callback = BraintrustDSpyCallback()
+            dspy.configure(lm=None, callbacks=[bt_callback])
+
+            from dspy.dsp.utils.settings import settings
+            callbacks = settings.callbacks
+
+            # Should only have one BraintrustDSpyCallback
+            bt_callbacks = [cb for cb in callbacks if isinstance(cb, BraintrustDSpyCallback)]
+            assert len(bt_callbacks) == 1, f"Expected 1 BraintrustDSpyCallback, got {len(bt_callbacks)}"
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+    def test_patch_dspy_idempotent(self):
+        """Multiple patch_dspy() calls should be safe."""
+        result = run_in_subprocess("""
+            from braintrust.wrappers.dspy import patch_dspy
+            import dspy
+
+            patch_dspy()
+            patch_dspy()  # Second call - should be no-op, not double-wrap
+
+            # Verify configure still works
+            lm = dspy.LM("openai/gpt-4o-mini")
+            dspy.configure(lm=lm)
+            print("SUCCESS")
+        """)
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+        assert "SUCCESS" in result.stdout
+
+
+class TestAutoInstrumentDSPy:
+    """Tests for auto_instrument() with DSPy."""
+
+    def test_auto_instrument_dspy(self):
+        """Test auto_instrument patches DSPy, creates spans, and uninstrument works."""
+        verify_autoinstrument_script("test_auto_dspy.py")

braintrust/wrappers/test_google_genai.py

@@ -6,6 +6,7 @@ import pytest
 from braintrust import logger
 from braintrust.test_helpers import init_test_logger
 from braintrust.wrappers.google_genai import setup_genai
+from braintrust.wrappers.test_utils import verify_autoinstrument_script
 from google.genai import types
 from google.genai.client import Client
 
@@ -637,3 +638,11 @@ def test_attachment_with_pydantic_model(memory_logger):
 
     # Attachment should be preserved
     assert copied["context_file"] is attachment
+
+
+class TestAutoInstrumentGoogleGenAI:
+    """Tests for auto_instrument() with Google GenAI."""
+
+    def test_auto_instrument_google_genai(self):
+        """Test auto_instrument patches Google GenAI and creates spans."""
+        verify_autoinstrument_script("test_auto_google_genai.py")