braintrust 0.5.0__py3-none-any.whl → 0.5.3__py3-none-any.whl

braintrust/test_merge_row_batch.py

@@ -0,0 +1,160 @@
+import unittest
+
+from braintrust.db_fields import IS_MERGE_FIELD
+from braintrust.merge_row_batch import batch_items, merge_row_batch
+
+
+class MergeRowBatchTest(unittest.TestCase):
+    def test_basic(self):
+        rows = [
+            # These rows should get merged together, ending up as a merge.
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(a=12),
+                **{IS_MERGE_FIELD: True},
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(b=10),
+                **{IS_MERGE_FIELD: True},
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(c="hello"),
+                **{IS_MERGE_FIELD: True},
+            ),
+            # The first row should be clobbered by the second, but the third
+            # merged with the second, ending up as a replacement.
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(a="hello"),
+            ),
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(b=10),
+            ),
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(c=12),
+                **{IS_MERGE_FIELD: True},
+            ),
+            # These rows should be clobbered separately from the last batch.
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(a="hello"),
+            ),
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(b=10),
+            ),
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(c=12),
+            ),
+        ]
+
+        merged_rows = merge_row_batch(rows)
+        key_to_rows = {(row.get("experiment_id"), row.get("dataset_id"), row.get("id")): row for row in merged_rows}
+        self.assertEqual(
+            {
+                ("e0", None, "x"): dict(
+                    experiment_id="e0",
+                    id="x",
+                    inputs=dict(a=12, b=10, c="hello"),
+                    **{IS_MERGE_FIELD: True},
+                ),
+                ("e0", None, "y"): dict(
+                    experiment_id="e0",
+                    id="y",
+                    inputs=dict(b=10, c=12),
+                ),
+                (None, "d0", "y"): dict(
+                    dataset_id="d0",
+                    id="y",
+                    inputs=dict(c=12),
+                ),
+            },
+            key_to_rows,
+        )
+
+    def test_skip_fields(self):
+        rows = [
+            # These rows should get merged together, ending up as a merge. But
+            # the original fields should be retained, regardless of whether we
+            # populated them or not.
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(a=12),
+                **{IS_MERGE_FIELD: True},
+                created=123,
+                root_span_id="abc",
+                _parent_id="baz",
+                span_parents=["foo", "bar"],
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(b=10),
+                **{IS_MERGE_FIELD: True},
+                created=456,
+                span_id="foo",
+                root_span_id="bar",
+                _parent_id="boop",
+                span_parents=[],
+            ),
+        ]
+
+        merged_rows = merge_row_batch(rows)
+        self.assertEqual(
+            merged_rows,
+            [
+                dict(
+                    experiment_id="e0",
+                    id="x",
+                    inputs=dict(a=12, b=10),
+                    **{IS_MERGE_FIELD: True},
+                    created=123,
+                    root_span_id="abc",
+                    _parent_id="baz",
+                    span_parents=["foo", "bar"],
+                ),
+            ],
+        )
+
+
+class BatchItemsTest(unittest.TestCase):
+    def test_basic(self):
+        a = "x" * 1
+        b = "x" * 2
+        c = "x" * 4
+        d = "y" * 1
+        e = "y" * 2
+        f = "y" * 4
+
+        items = [a, b, c, f, e, d]
+
+        # No limits.
+        output = batch_items(items)
+        self.assertEqual(output, [[a, b, c, f, e, d]])
+
+        # Num items limit.
+        output = batch_items(items, batch_max_num_items=2)
+        self.assertEqual(output, [[a, b], [c, f], [e, d]])
+
+        # Num bytes limit.
+        output = batch_items(items, batch_max_num_bytes=2)
+        self.assertEqual(output, [[a], [b], [c], [f], [e], [d]])
+
+        # Both items and num bytes limit.
+        output = batch_items(items, batch_max_num_items=2, batch_max_num_bytes=5)
+        self.assertEqual(output, [[a, b], [c], [f], [e, d]])

braintrust/test_util.py

@@ -1,9 +1,66 @@
+import os
 import unittest
 from typing import List
 
 import pytest
 
-from .util import LazyValue, mask_api_key, merge_dicts_with_paths
+from .util import LazyValue, mask_api_key, merge_dicts_with_paths, parse_env_var_float
+
+
+class TestParseEnvVarFloat:
+    """Tests for parse_env_var_float helper."""
+
+    def test_returns_default_when_env_not_set(self):
+        assert parse_env_var_float("NONEXISTENT_VAR_12345", 42.0) == 42.0
+
+    def test_parses_valid_float(self):
+        os.environ["TEST_FLOAT"] = "123.45"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 0.0) == 123.45
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_returns_default_for_nan(self):
+        os.environ["TEST_FLOAT"] = "nan"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_returns_default_for_inf(self):
+        os.environ["TEST_FLOAT"] = "inf"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_returns_default_for_negative_inf(self):
+        os.environ["TEST_FLOAT"] = "-inf"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_returns_default_for_empty_string(self):
+        os.environ["TEST_FLOAT"] = ""
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_returns_default_for_invalid_string(self):
+        os.environ["TEST_FLOAT"] = "not_a_number"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 99.0) == 99.0
+        finally:
+            del os.environ["TEST_FLOAT"]
+
+    def test_allows_negative_values(self):
+        os.environ["TEST_FLOAT"] = "-5.5"
+        try:
+            assert parse_env_var_float("TEST_FLOAT", 0.0) == -5.5
+        finally:
+            del os.environ["TEST_FLOAT"]
 
 
 class TestLazyValue(unittest.TestCase):

braintrust/util.py

@@ -1,5 +1,7 @@
 import inspect
 import json
+import math
+import os
 import sys
 import threading
 import urllib.parse
@@ -9,6 +11,24 @@ from typing import Any, Generic, Literal, TypedDict, TypeVar, Union
 
 from requests import HTTPError, Response
 
+
+def parse_env_var_float(name: str, default: float) -> float:
+    """Parse a float from an environment variable, returning default if invalid.
+
+    Returns the default value if the env var is missing, empty, not a valid
+    float, NaN, or infinity.
+    """
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    try:
+        result = float(value)
+        if math.isnan(result) or math.isinf(result):
+            return default
+        return result
+    except (ValueError, TypeError):
+        return default
+
 GLOBAL_PROJECT = "Global"
 BT_IS_ASYNC_ATTRIBUTE = "_BT_IS_ASYNC"
 

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.5.0"
+VERSION = "0.5.3"
 
 # this will be templated during the build
-GIT_COMMIT = "617d9b730b37e96b7d05a099b95f5387944d0951"
+GIT_COMMIT = "ce261c13723e8985e61819f6df56298f863155d0"

braintrust/wrappers/agno/__init__.py

@@ -62,7 +62,6 @@ def setup_agno(
         models.base.Model = wrap_model(models.base.Model)  # pyright: ignore[reportUnknownMemberType]
         tools.function.FunctionCall = wrap_function_call(tools.function.FunctionCall)  # pyright: ignore[reportUnknownMemberType]
         return True
-    except ImportError as e:
-        logger.error(f"Failed to import Agno: {e}")
-        logger.error("Agno is not installed. Please install it with: pip install agno")
+    except ImportError:
+        # Not installed - this is expected when using auto_instrument()
         return False

braintrust/wrappers/anthropic.py

@@ -5,6 +5,7 @@ from contextlib import contextmanager
 
 from braintrust.logger import NOOP_SPAN, log_exc_info_to_span, start_span
 from braintrust.wrappers._anthropic_utils import Wrapper, extract_anthropic_usage, finalize_anthropic_tokens
+from wrapt import wrap_function_wrapper
 
 log = logging.getLogger(__name__)
 
@@ -358,3 +359,66 @@ def wrap_anthropic(client):
 
 def wrap_anthropic_client(client):
     return wrap_anthropic(client)
+
+
+def _apply_anthropic_wrapper(client):
+    """Apply tracing wrapper to an Anthropic client instance in-place."""
+    wrapped = wrap_anthropic(client)
+    client.messages = wrapped.messages
+    if hasattr(wrapped, "beta"):
+        client.beta = wrapped.beta
+
+
+def _apply_async_anthropic_wrapper(client):
+    """Apply tracing wrapper to an AsyncAnthropic client instance in-place."""
+    wrapped = wrap_anthropic(client)
+    client.messages = wrapped.messages
+    if hasattr(wrapped, "beta"):
+        client.beta = wrapped.beta
+
+
+def _anthropic_init_wrapper(wrapped, instance, args, kwargs):
+    """Wrapper for Anthropic.__init__ that applies tracing after initialization."""
+    wrapped(*args, **kwargs)
+    _apply_anthropic_wrapper(instance)
+
+
+def _async_anthropic_init_wrapper(wrapped, instance, args, kwargs):
+    """Wrapper for AsyncAnthropic.__init__ that applies tracing after initialization."""
+    wrapped(*args, **kwargs)
+    _apply_async_anthropic_wrapper(instance)
+
+
+def patch_anthropic() -> bool:
+    """
+    Patch Anthropic to add Braintrust tracing globally.
+
+    After calling this, all new Anthropic() and AsyncAnthropic() clients
+    will automatically have tracing enabled.
+
+    Returns:
+        True if Anthropic was patched (or already patched), False if Anthropic is not installed.
+
+    Example:
+        ```python
+        import braintrust
+        braintrust.patch_anthropic()
+
+        import anthropic
+        client = anthropic.Anthropic()
+        # All calls are now traced!
+        ```
+    """
+    try:
+        import anthropic
+
+        if getattr(anthropic, "__braintrust_wrapped__", False):
+            return True  # Already patched
+
+        wrap_function_wrapper("anthropic", "Anthropic.__init__", _anthropic_init_wrapper)
+        wrap_function_wrapper("anthropic", "AsyncAnthropic.__init__", _async_anthropic_init_wrapper)
+        anthropic.__braintrust_wrapped__ = True
+        return True
+
+    except ImportError:
+        return False

braintrust/wrappers/claude_agent_sdk/__init__.py

@@ -105,7 +105,6 @@ def setup_claude_agent_sdk(
                         setattr(module, "tool", wrapped_tool_fn)
 
         return True
-    except ImportError as e:
-        logger.error(f"Failed to import Claude Agent SDK: {e}")
-        logger.error("claude-agent-sdk is not installed. Please install it with: pip install claude-agent-sdk")
+    except ImportError:
+        # Not installed - this is expected when using auto_instrument()
         return False

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"
@@ -283,3 +284,11 @@ 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")