braintrust 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

braintrust/test_otel.py

@@ -294,13 +294,12 @@ class TestSpanFiltering:
             self.provider.shutdown()
             self.memory_exporter.clear()
 
-    def test_keeps_root_spans(self):
+    def test_filters_out_root_spans(self):
         with self.tracer.start_as_current_span("root_operation"):
             pass
 
         spans = self.memory_exporter.get_finished_spans()
-        assert len(spans) == 1
-        assert spans[0].name == "root_operation"
+        assert len(spans) == 0
 
     def test_keeps_gen_ai_spans(self):
         with self.tracer.start_as_current_span("root"):
@@ -312,7 +311,7 @@ class TestSpanFiltering:
         spans = self.memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
 
-        assert "root" in span_names
+        assert "root" not in span_names
         assert "gen_ai.completion" in span_names
         assert "regular_operation" not in span_names
 
@@ -329,35 +328,37 @@ class TestSpanFiltering:
         assert "braintrust.eval" in span_names
         assert "database_query" not in span_names
 
-    def test_keeps_llm_spans(self):
+    def test_keeps_traceloop_spans(self):
         with self.tracer.start_as_current_span("root"):
-            with self.tracer.start_as_current_span("llm.generate"):
+            with self.tracer.start_as_current_span("traceloop.agent"):
+                pass
+            with self.tracer.start_as_current_span("traceloop.workflow.step"):
                 pass
 
         spans = self.memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
-        assert "llm.generate" in span_names
+        assert "root" not in span_names
+        assert "traceloop.agent" in span_names
+        assert "traceloop.workflow.step" in span_names
 
-    def test_keeps_ai_spans(self):
+    def test_keeps_llm_spans(self):
         with self.tracer.start_as_current_span("root"):
-            with self.tracer.start_as_current_span("ai.model_call"):
+            with self.tracer.start_as_current_span("llm.generate"):
                 pass
 
         spans = self.memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
-        assert "ai.model_call" in span_names
+        assert "llm.generate" in span_names
 
-    def test_keeps_traceloop_spans(self):
+    def test_keeps_ai_spans(self):
         with self.tracer.start_as_current_span("root"):
-            with self.tracer.start_as_current_span("traceloop.agent"):
-                pass
-            with self.tracer.start_as_current_span("traceloop.workflow.step"):
+            with self.tracer.start_as_current_span("ai.model_call"):
                 pass
 
         spans = self.memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
-        assert "traceloop.agent" in span_names
-        assert "traceloop.workflow.step" in span_names
+        assert "root" not in span_names
+        assert "ai.model_call" in span_names
 
     def test_keeps_spans_with_llm_attributes(self):
         with self.tracer.start_as_current_span("root"):
@@ -374,7 +375,7 @@ class TestSpanFiltering:
         spans = self.memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
 
-        assert "root" in span_names
+        assert "root" not in span_names
         assert "some_operation" in span_names  # has gen_ai.model attribute
         assert "another_operation" in span_names  # has llm.tokens attribute
         assert "traceloop_operation" in span_names  # has traceloop.agent_id attribute
@@ -390,10 +391,7 @@ class TestSpanFiltering:
                 pass
 
         spans = self.memory_exporter.get_finished_spans()
-
-        # Only root should be kept
-        assert len(spans) == 1
-        assert spans[0].name == "root"
+        assert len(spans) == 0
 
     def test_custom_filter_keeps_spans(self):
         def custom_filter(span):
@@ -422,9 +420,9 @@ class TestSpanFiltering:
         spans = memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
 
-        assert "root" in span_names
         assert "custom_keep" in span_names  # kept by custom filter
         assert "regular_operation" not in span_names  # dropped by default logic
+        assert "root" not in span_names
 
     def test_custom_filter_drops_spans(self):
         def custom_filter(span):
@@ -453,9 +451,9 @@ class TestSpanFiltering:
         spans = memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
 
-        assert "root" in span_names
         assert "gen_ai.drop_this" not in span_names  # dropped by custom filter
         assert "gen_ai.keep_this" in span_names  # kept by default LLM logic
+        assert "root" not in span_names
 
     def test_custom_filter_none_uses_default_logic(self):
         def custom_filter(span):
@@ -482,7 +480,7 @@ class TestSpanFiltering:
         spans = memory_exporter.get_finished_spans()
         span_names = [span.name for span in spans]
 
-        assert "root" in span_names
+        assert "root" not in span_names
         assert "gen_ai.completion" in span_names  # kept by default LLM logic
         assert "regular_operation" not in span_names  # dropped by default logic
 
@@ -546,11 +544,32 @@ class TestSpanFiltering:
         filtered_spans = filtered_spans_exporter.get_finished_spans()
         filtered_span_names = [span.name for span in filtered_spans]
 
-        assert len(filtered_spans) == 3
-        assert "user_request" in filtered_span_names  # root span
+        assert len(filtered_spans) == 2
+        assert "user_request" not in filtered_span_names  # root span
         assert "gen_ai.completion" in filtered_span_names  # LLM name
         assert "response_formatting" in filtered_span_names  # LLM attribute
 
+    def test_custom_filter_is_root_span(self):
+        from braintrust.otel import AISpanProcessor, is_root_span
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+        from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+        memory_exporter = InMemorySpanExporter()
+        processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=is_root_span)
+        provider = TracerProvider()
+        provider.add_span_processor(processor)
+        tracer = provider.get_tracer("test-braintrust-root-filter")
+
+        with tracer.start_as_current_span("root_span"):
+            with tracer.start_as_current_span("child_span"):
+                pass
+
+        provider.shutdown()
+        spans = memory_exporter.get_finished_spans()
+        names = [span.name for span in spans]
+        assert "root_span" in names
+        assert "child_span" not in names
 
 def test_parent_from_headers_invalid_inputs():
     """Test parent_from_headers with various invalid inputs."""
@@ -716,3 +735,76 @@ def test_add_span_parent_to_baggage():
     # Test with None span (should return None and warn)
     token = add_span_parent_to_baggage(None)
     assert token is None
+
+
+def test_parent_from_headers_with_custom_propagator():
+    """Test parent_from_headers with a custom propagator."""
+    if not _check_otel_installed():
+        pytest.skip("OpenTelemetry SDK not fully installed, skipping test")
+
+    from braintrust.otel import parent_from_headers
+    from opentelemetry import baggage as otel_baggage
+    from opentelemetry import context as otel_context
+    from opentelemetry import trace
+    from opentelemetry.propagators.textmap import CarrierT, Getter, TextMapPropagator, default_getter
+    from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
+
+    class CustomHeaderPropagator(TextMapPropagator):
+        """Custom propagator that reads trace context from X-Custom-* headers."""
+
+        def extract(
+            self,
+            carrier: CarrierT,
+            context: otel_context.Context | None = None,
+            getter: Getter = default_getter,
+        ) -> otel_context.Context:
+            if context is None:
+                context = otel_context.get_current()
+
+            trace_id = getter.get(carrier, "X-Custom-Trace-Id")
+            span_id = getter.get(carrier, "X-Custom-Span-Id")
+
+            if trace_id and span_id:
+                trace_id_list = trace_id if isinstance(trace_id, list) else [trace_id]
+                span_id_list = span_id if isinstance(span_id, list) else [span_id]
+
+                span_context = SpanContext(
+                    trace_id=int(trace_id_list[0], 16),
+                    span_id=int(span_id_list[0], 16),
+                    is_remote=True,
+                    trace_flags=TraceFlags.SAMPLED,
+                )
+                span = NonRecordingSpan(span_context)
+                context = trace.set_span_in_context(span, context)
+
+            # Also extract baggage from standard baggage header
+            baggage_header = getter.get(carrier, "baggage")
+            if baggage_header:
+                baggage_list = baggage_header if isinstance(baggage_header, list) else [baggage_header]
+                for item in baggage_list[0].split(","):
+                    if "=" in item:
+                        key, value = item.split("=", 1)
+                        context = otel_baggage.set_baggage(key.strip(), value.strip(), context)
+
+            return context
+
+        def inject(self, carrier, context=None, setter=None):
+            pass  # Not needed for this test
+
+        @property
+        def fields(self):
+            return {"X-Custom-Trace-Id", "X-Custom-Span-Id", "baggage"}
+
+    propagator = CustomHeaderPropagator()
+
+    # Custom header format
+    headers = {
+        "X-Custom-Trace-Id": "4bf92f3577b34da6a3ce929d0e0e4736",
+        "X-Custom-Span-Id": "00f067aa0ba902b7",
+        "baggage": "braintrust.parent=project_name:test-project",
+    }
+
+    result = parent_from_headers(headers, propagator=propagator)
+    assert result is not None
+    assert isinstance(result, str)
+    assert len(result) > 0

braintrust/test_util.py

@@ -3,7 +3,7 @@ from typing import List
 
 import pytest
 
-from .util import LazyValue, mask_api_key
+from .util import LazyValue, mask_api_key, merge_dicts_with_paths
 
 
 class TestLazyValue(unittest.TestCase):
@@ -160,3 +160,53 @@ def test_mask_api_key():
     assert mask_api_key("12345") == "12*45"
     for i in ["", "1", "12", "123", "1234"]:
         assert mask_api_key(i) == "*" * len(i)
+
+
+class TestTagsSetUnionMerge:
+    def test_tags_arrays_are_merged_as_sets_by_default(self):
+        a = {"tags": ["a", "b"]}
+        b = {"tags": ["b", "c"]}
+        merge_dicts_with_paths(a, b, (), set())
+        assert set(a["tags"]) == {"a", "b", "c"}
+
+    def test_tags_merge_deduplicates_values(self):
+        a = {"tags": ["a", "b", "c"]}
+        b = {"tags": ["a", "b", "c", "d"]}
+        merge_dicts_with_paths(a, b, (), set())
+        assert set(a["tags"]) == {"a", "b", "c", "d"}
+
+    def test_tags_merge_works_when_merge_into_has_no_tags(self):
+        a = {"other": "data"}
+        b = {"tags": ["a", "b"]}
+        merge_dicts_with_paths(a, b, (), set())
+        assert set(a["tags"]) == {"a", "b"}
+
+    def test_tags_merge_works_when_merge_from_has_no_tags(self):
+        a = {"tags": ["a", "b"]}
+        b = {"other": "data"}
+        merge_dicts_with_paths(a, b, (), set())
+        assert set(a["tags"]) == {"a", "b"}
+
+    def test_tags_are_replaced_when_included_in_merge_paths(self):
+        a = {"tags": ["a", "b"]}
+        b = {"tags": ["c", "d"]}
+        merge_dicts_with_paths(a, b, (), {("tags",)})
+        assert a["tags"] == ["c", "d"]
+
+    def test_empty_tags_array_clears_tags_when_in_merge_paths(self):
+        a = {"tags": ["a", "b"]}
+        b = {"tags": []}
+        merge_dicts_with_paths(a, b, (), {("tags",)})
+        assert a["tags"] == []
+
+    def test_none_tags_replaces_tags(self):
+        a = {"tags": ["a", "b"]}
+        b = {"tags": None}
+        merge_dicts_with_paths(a, b, (), set())
+        assert a["tags"] is None
+
+    def test_set_union_only_applies_to_top_level_tags_field(self):
+        a = {"metadata": {"tags": ["a", "b"]}}
+        b = {"metadata": {"tags": ["c", "d"]}}
+        merge_dicts_with_paths(a, b, (), set())
+        assert a["metadata"]["tags"] == ["c", "d"]

braintrust/util.py

@@ -1,4 +1,5 @@
 import inspect
+import json
 import sys
 import threading
 import urllib.parse
@@ -29,11 +30,16 @@ def coalesce(*args):
     return None
 
 
+# Fields that automatically use set-union merge semantics (unless in merge_paths).
+_SET_UNION_FIELDS = frozenset(["tags"])
+
+
 def merge_dicts_with_paths(
-    merge_into: dict[str, Any], merge_from: Mapping[str, Any], path: tuple[str, ...], merge_paths: set[tuple[str]]
+    merge_into: dict[str, Any], merge_from: Mapping[str, Any], path: tuple[str, ...], merge_paths: set[tuple[str, ...]]
 ) -> dict[str, Any]:
     """Merges merge_from into merge_into, destructively updating merge_into. Does not merge any further than
-    merge_paths."""
+    merge_paths. For fields in _SET_UNION_FIELDS (like "tags"), arrays are merged as sets (union)
+    unless the field is explicitly listed in merge_paths (opt-out to replacement)."""
 
     if not isinstance(merge_into, dict):
         raise ValueError("merge_into must be a dictionary")
@@ -43,7 +49,22 @@ def merge_dicts_with_paths(
     for k, merge_from_v in merge_from.items():
         full_path = path + (k,)
         merge_into_v = merge_into.get(k)
-        if isinstance(merge_into_v, dict) and isinstance(merge_from_v, dict) and full_path not in merge_paths:
+
+        # Check if this field should use set-union merge (e.g., "tags" at top level)
+        is_set_union_field = len(path) == 0 and k in _SET_UNION_FIELDS and full_path not in merge_paths
+
+        if is_set_union_field and isinstance(merge_into_v, list) and isinstance(merge_from_v, list):
+            # Set-union merge: combine arrays, deduplicate using JSON for objects
+            seen: set[str] = set()
+            combined = []
+            for item in merge_into_v + list(merge_from_v):
+                # Use JSON serialization for consistent object comparison
+                item_key = json.dumps(item, sort_keys=True) if isinstance(item, (dict, list)) else str(item)
+                if item_key not in seen:
+                    seen.add(item_key)
+                    combined.append(item)
+            merge_into[k] = combined
+        elif isinstance(merge_into_v, dict) and isinstance(merge_from_v, dict) and full_path not in merge_paths:
             merge_dicts_with_paths(merge_into_v, merge_from_v, full_path, merge_paths)
         else:
             merge_into[k] = merge_from_v

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.4.0"
+VERSION = "0.4.2"
 
 # this will be templated during the build
-GIT_COMMIT = "8ab13f3f48af6a4d3c0b053e4bbabfd4f24f23ec"
+GIT_COMMIT = "3ca420e53e77d4665b91ccc7631c95dc97ce566d"

braintrust/wrappers/google_genai/__init__.py

@@ -3,6 +3,7 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+from braintrust.bt_json import bt_safe_deep_copy
 from braintrust.logger import NOOP_SPAN, Attachment, current_span, init_logger, start_span
 from braintrust.span_types import SpanTypeAttribute
 from wrapt import wrap_function_wrapper
@@ -149,7 +150,7 @@ def wrap_async_models(AsyncModels: Any):
 
 
 def _serialize_input(api_client: Any, input: dict[str, Any]):
-    config = _try_dict(input.get("config"))
+    config = bt_safe_deep_copy(input.get("config"))
 
     if config is not None:
         tools = _serialize_tools(api_client, input)
@@ -424,17 +425,3 @@ def get_path(obj: dict[str, Any], path: str, default: Any = None) -> Any | None:
         current = current[key]
 
     return current
-
-
-def _try_dict(obj: Any) -> dict[str, Any] | None:
-    try:
-        return obj.model_dump()
-    except AttributeError:
-        pass
-
-    try:
-        return obj.dump()
-    except AttributeError:
-        pass
-
-    return obj

braintrust/wrappers/litellm.py

@@ -657,9 +657,52 @@ 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
     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