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

braintrust/span_identifier_v4.py

@@ -5,7 +5,6 @@ import base64
 import dataclasses
 import json
 from enum import Enum
-from typing import Dict, Optional, Union
 
 from .span_identifier_v3 import (
     SpanComponentsV3,
@@ -14,6 +13,7 @@ from .span_identifier_v3 import (
 
 ENCODING_VERSION_NUMBER_V4 = 4
 
+
 def _try_make_hex_trace_id(s):
     """Try to convert hex string to 16-byte binary (for trace IDs)"""
     try:
@@ -25,6 +25,7 @@ def _try_make_hex_trace_id(s):
         pass
     return None, False
 
+
 def _try_make_hex_span_id(s):
     """Try to convert hex string to 8-byte binary (for span IDs)"""
     try:
@@ -36,14 +37,17 @@ def _try_make_hex_span_id(s):
         pass
     return None, False
 
+
 INVALID_ENCODING_ERRMSG_V4 = f"SpanComponents string is not properly encoded. This library only supports encoding versions up to {ENCODING_VERSION_NUMBER_V4}. Please make sure the SDK library used to decode the SpanComponents is at least as new as any library used to encode it."
 
+
 class Fields(Enum):
     OBJECT_ID = 1
     ROW_ID = 2
     SPAN_ID = 3  # 8-byte hex
     ROOT_SPAN_ID = 4  # 16-byte hex
 
+
 _FIELDS_ID_TO_NAME = {
     Fields.OBJECT_ID: "object_id",
     Fields.ROW_ID: "row_id",
@@ -57,16 +61,16 @@ class SpanComponentsV4:
     object_type: SpanObjectTypeV3
 
     # Must provide one or the other.
-    object_id: Optional[str] = None
-    compute_object_metadata_args: Optional[Dict] = None
+    object_id: str | None = None
+    compute_object_metadata_args: dict | None = None
 
     # Either all of these must be provided or none.
-    row_id: Optional[str] = None
-    span_id: Optional[str] = None
-    root_span_id: Optional[str] = None
+    row_id: str | None = None
+    span_id: str | None = None
+    root_span_id: str | None = None
 
     # Additional span properties.
-    propagated_event: Optional[Dict] = None
+    propagated_event: dict | None = None
 
     def __post_init__(self):
         # Reuse V3 validation logic
@@ -98,10 +102,12 @@ class SpanComponentsV4:
         )
         json_obj = {k: v for k, v in json_obj.items() if v is not None}
 
-        raw_bytes = bytes([
-            ENCODING_VERSION_NUMBER_V4,
-            self.object_type.value,
-        ])
+        raw_bytes = bytes(
+            [
+                ENCODING_VERSION_NUMBER_V4,
+                self.object_type.value,
+            ]
+        )
 
         hex_entries = []
 
@@ -199,7 +205,7 @@ class SpanComponentsV4:
         except Exception:
             raise Exception(INVALID_ENCODING_ERRMSG_V4)
 
-    def object_id_fields(self) -> Dict[str, str]:
+    def object_id_fields(self) -> dict[str, str]:
         # Reuse V3 logic
         if not self.object_id:
             raise Exception(
@@ -218,7 +224,7 @@ class SpanComponentsV4:
         return self.to_str()
 
     @staticmethod
-    def _from_json_obj(json_obj: Dict) -> "SpanComponentsV4":
+    def _from_json_obj(json_obj: dict) -> "SpanComponentsV4":
         kwargs = {
             **json_obj,
             "object_type": SpanObjectTypeV3(json_obj["object_type"]),
@@ -226,8 +232,7 @@ class SpanComponentsV4:
         return SpanComponentsV4(**kwargs)
 
 
-
-def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
+def parse_parent(parent: str | dict | None) -> str | None:
     """Parse a parent object into a string representation using V4 format."""
     # Reuse V3 logic but with V4 components
     if isinstance(parent, str):
@@ -250,17 +255,21 @@ def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
 
         row_ids = parent.get("row_ids")
         if row_ids:
-            kwargs.update({
-                "row_id": row_ids.get("id"),
-                "span_id": row_ids.get("span_id"),
-                "root_span_id": row_ids.get("root_span_id"),
-            })
+            kwargs.update(
+                {
+                    "row_id": row_ids.get("id"),
+                    "span_id": row_ids.get("span_id"),
+                    "root_span_id": row_ids.get("root_span_id"),
+                }
+            )
         else:
-            kwargs.update({
-                "row_id": None,
-                "span_id": None,
-                "root_span_id": None,
-            })
+            kwargs.update(
+                {
+                    "row_id": None,
+                    "span_id": None,
+                    "root_span_id": None,
+                }
+            )
 
         if "propagated_event" in parent:
             kwargs["propagated_event"] = parent.get("propagated_event")

braintrust/test_framework.py

@@ -240,12 +240,15 @@ async def test_hooks_trial_index_multiple_inputs():
     assert sorted(input_1_trials) == [0, 1]
     assert sorted(input_2_trials) == [0, 1]
 
+
 @pytest.fixture
 def simple_scorer():
     def simple_scorer_function(input, output, expected):
         return {"name": "simple_scorer", "score": 0.8}
+
     return simple_scorer_function
 
+
 @pytest.mark.asyncio
 async def test_eval_no_send_logs_true(with_memory_logger, simple_scorer):
     """Test that Eval with no_send_logs=True runs locally without creating experiment."""
@@ -286,7 +289,7 @@ async def test_eval_no_send_logs_true(with_memory_logger, simple_scorer):
 
 @pytest.mark.asyncio
 async def test_hooks_tags_append(with_memory_logger, with_simulate_login, simple_scorer):
-    """ Test that hooks.tags can be appended to and logged. """
+    """Test that hooks.tags can be appended to and logged."""
 
     initial_tags = ["cookies n cream"]
     appended_tags = ["chocolate", "vanilla", "strawberry"]
@@ -321,9 +324,12 @@ async def test_hooks_tags_append(with_memory_logger, with_simulate_login, simple
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize(("tags", "expected_tags"), [(None, None),([], None), (["chocolate", "vanilla", "strawberry"], ["chocolate", "vanilla", "strawberry"])])
+@pytest.mark.parametrize(
+    ("tags", "expected_tags"),
+    [(None, None), ([], None), (["chocolate", "vanilla", "strawberry"], ["chocolate", "vanilla", "strawberry"])],
+)
 async def test_hooks_tags_list(with_memory_logger, with_simulate_login, simple_scorer, tags, expected_tags):
-    """ Test that hooks.tags can be set to a list. """
+    """Test that hooks.tags can be set to a list."""
 
     def task_with_hooks(input, hooks):
         hooks.tags = tags
@@ -351,9 +357,10 @@ async def test_hooks_tags_list(with_memory_logger, with_simulate_login, simple_s
     assert len(root_span) == 1
     assert root_span[0].get("tags") == expected_tags
 
+
 @pytest.mark.asyncio
 async def test_hooks_tags_with_failing_scorer(with_memory_logger, with_simulate_login, simple_scorer):
-    """ Test that hooks.tags can be set to a list. """
+    """Test that hooks.tags can be set to a list."""
 
     expected_tags = ["chocolate", "vanilla", "strawberry"]
 
@@ -386,9 +393,11 @@ async def test_hooks_tags_with_failing_scorer(with_memory_logger, with_simulate_
     assert len(root_span) == 1
     assert root_span[0].get("tags") == expected_tags
 
+
 @pytest.mark.asyncio
 async def test_hooks_tags_with_invalid_type(with_memory_logger, with_simulate_login, simple_scorer):
-    """ Test that result contains an error for cases where hooks.tags is set to an invalid type. """
+    """Test that result contains an error for cases where hooks.tags is set to an invalid type."""
+
     def task_with_hooks(input, hooks):
         hooks.tags = 123
         return input
@@ -411,7 +420,8 @@ async def test_hooks_tags_with_invalid_type(with_memory_logger, with_simulate_lo
 
 @pytest.mark.asyncio
 async def test_hooks_without_setting_tags(with_memory_logger, with_simulate_login, simple_scorer):
-    """ Test where hooks.tags is not set """
+    """Test where hooks.tags is not set"""
+
     def task_with_hooks(input, hooks):
         return input
 

braintrust/test_helpers.py

@@ -2,7 +2,6 @@ import os
 from contextlib import contextmanager
 
 import pytest
-
 from braintrust import logger
 from braintrust.logger import ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata
 from braintrust.util import LazyValue
@@ -15,10 +14,8 @@ TEST_ORG_NAME = "test-org-name"
 def has_devserver_installed() -> bool:
     """Check if devserver dependencies (starlette, uvicorn) are installed."""
     import importlib.util
-    return (
-        importlib.util.find_spec("starlette") is not None
-        and importlib.util.find_spec("uvicorn") is not None
-    )
+
+    return importlib.util.find_spec("starlette") is not None and importlib.util.find_spec("uvicorn") is not None
 
 
 def simulate_login() -> None:
@@ -68,12 +65,14 @@ def with_memory_logger():
     # Clean up global state to prevent test contamination
     logger._state.reset_parent_state()
 
+
 @pytest.fixture
 def memory_logger():
     with logger._internal_with_memory_background_logger() as bgl:
         yield bgl
     logger._state.current_experiment = None
 
+
 @contextmanager
 def preserve_env_vars(*vars):
     original_env = {v: os.environ.get(v) for v in vars}
@@ -114,6 +113,7 @@ def init_test_logger(project_name: str):
     logger._compute_logger_metadata = fake_compute_logger_metadata
     return l
 
+
 def init_test_exp(experiment_name: str, project_name: str = None):
     """
     Initialize an experiment for testing with fake project and experiment metadata.

braintrust/test_id_gen.py

@@ -1,9 +1,7 @@
-
 import os
 import uuid
 
 import pytest
-
 from braintrust import id_gen
 
 
@@ -76,7 +74,8 @@ def test_id_get_env_var(reset_id_generator_state):
 
 
 def _is_hex(s):
-    return all(c in '0123456789abcdef' for c in s.lower())
+    return all(c in "0123456789abcdef" for c in s.lower())
+
 
 def _assert_is_hex(x):
     assert _is_hex(x)

braintrust/test_otel.py

@@ -274,12 +274,11 @@ class TestSpanFiltering:
         except ImportError:
             pytest.skip("OpenTelemetry SDK not fully installed, skipping AISpanProcessor tests")
 
+        from braintrust.otel import AISpanProcessor
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor
         from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-        from braintrust.otel import AISpanProcessor
-
         self.memory_exporter = InMemorySpanExporter()
         self.provider = TracerProvider()
 
@@ -403,12 +402,11 @@ class TestSpanFiltering:
             return None  # Don't influence decision
 
         # Create processor with custom filter
+        from braintrust.otel import AISpanProcessor
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor
         from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-        from braintrust.otel import AISpanProcessor
-
         memory_exporter = InMemorySpanExporter()
         processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
         provider = TracerProvider()
@@ -435,12 +433,11 @@ class TestSpanFiltering:
             return None  # Don't influence decision
 
         # Create processor with custom filter
+        from braintrust.otel import AISpanProcessor
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor
         from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-        from braintrust.otel import AISpanProcessor
-
         memory_exporter = InMemorySpanExporter()
         processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
         provider = TracerProvider()
@@ -465,12 +462,11 @@ class TestSpanFiltering:
             return None  # Always defer to default logic
 
         # Create processor with custom filter
+        from braintrust.otel import AISpanProcessor
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor
         from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-        from braintrust.otel import AISpanProcessor
-
         memory_exporter = InMemorySpanExporter()
         processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
         provider = TracerProvider()
@@ -492,12 +488,11 @@ class TestSpanFiltering:
 
     def test_filtering_vs_unfiltered_comparison(self):
         # Set up two separate exporters and processors
+        from braintrust.otel import AISpanProcessor
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor
         from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-        from braintrust.otel import AISpanProcessor
-
         all_spans_exporter = InMemorySpanExporter()
         filtered_spans_exporter = InMemorySpanExporter()
 
@@ -569,49 +564,58 @@ def test_parent_from_headers_invalid_inputs():
     assert result is None
 
     # Test 2: Invalid traceparent (malformed)
-    result = parent_from_headers({'traceparent': 'invalid'})
+    result = parent_from_headers({"traceparent": "invalid"})
     assert result is None
 
     # Test 3: Valid traceparent but invalid braintrust.parent format
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=invalid_format'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=invalid_format",
+        }
+    )
     assert result is None
 
     # Test 4: Empty project_id
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=project_id:'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=project_id:",
+        }
+    )
     assert result is None
 
     # Test 5: Empty project_name
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=project_name:'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=project_name:",
+        }
+    )
     assert result is None
 
     # Test 6: Empty experiment_id
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=experiment_id:'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=experiment_id:",
+        }
+    )
     assert result is None
 
     # Test 7: Invalid trace_id length (too short)
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=project_name:test'
-    })
+    result = parent_from_headers(
+        {"traceparent": "00-4bf92f3577b34da6-00f067aa0ba902b7-01", "baggage": "braintrust.parent=project_name:test"}
+    )
     assert result is None
 
     # Test 8: Invalid span_id length (too short)
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa-01',
-        'baggage': 'braintrust.parent=project_name:test'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa-01",
+            "baggage": "braintrust.parent=project_name:test",
+        }
+    )
     assert result is None
 
 
@@ -623,29 +627,35 @@ def test_parent_from_headers_valid_input():
     from braintrust.otel import parent_from_headers
 
     # Test with valid project_name
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=project_name:test-project'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=project_name:test-project",
+        }
+    )
     assert result is not None
     # Result is base64 encoded, so just check it's a non-empty string
     assert isinstance(result, str)
     assert len(result) > 0
 
     # Test with valid project_id
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=project_id:abc123'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=project_id:abc123",
+        }
+    )
     assert result is not None
     assert isinstance(result, str)
     assert len(result) > 0
 
     # Test with valid experiment_id
-    result = parent_from_headers({
-        'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
-        'baggage': 'braintrust.parent=experiment_id:exp-456'
-    })
+    result = parent_from_headers(
+        {
+            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+            "baggage": "braintrust.parent=experiment_id:exp-456",
+        }
+    )
     assert result is not None
     assert isinstance(result, str)
     assert len(result) > 0
@@ -656,16 +666,15 @@ def test_add_parent_to_baggage():
     if not _check_otel_installed():
         pytest.skip("OpenTelemetry SDK not fully installed, skipping test")
 
-    from opentelemetry import baggage, context
-
     from braintrust.otel import add_parent_to_baggage
+    from opentelemetry import baggage, context
 
     # Test adding parent to baggage
     token = add_parent_to_baggage("project_name:test-project")
     assert token is not None
 
     # Verify it's in baggage
-    parent_value = baggage.get_baggage('braintrust.parent')
+    parent_value = baggage.get_baggage("braintrust.parent")
     assert parent_value == "project_name:test-project"
 
     # Clean up
@@ -677,11 +686,10 @@ def test_add_span_parent_to_baggage():
     if not _check_otel_installed():
         pytest.skip("OpenTelemetry SDK not fully installed, skipping test")
 
+    from braintrust.otel import add_span_parent_to_baggage
     from opentelemetry import baggage, context, trace
     from opentelemetry.sdk.trace import TracerProvider
 
-    from braintrust.otel import add_span_parent_to_baggage
-
     # Setup tracer
     provider = TracerProvider()
     trace.set_tracer_provider(provider)
@@ -695,7 +703,7 @@ def test_add_span_parent_to_baggage():
         assert token is not None
 
         # Verify it's in baggage
-        parent_value = baggage.get_baggage('braintrust.parent')
+        parent_value = baggage.get_baggage("braintrust.parent")
         assert parent_value == "project_name:test"
 
         context.detach(token)

braintrust/test_queue.py

@@ -3,7 +3,6 @@ import threading
 import time
 
 import pytest
-
 from braintrust.queue import DEFAULT_QUEUE_SIZE, LogQueue
 
 

braintrust/test_score.py

@@ -65,9 +65,7 @@ class TestScore(unittest.TestCase):
 
     def test_from_dict_round_trip(self):
         """Test that Score can be serialized to dict and deserialized back."""
-        original = Score(
-            name="round_trip_scorer", score=0.95, metadata={"info": "test"}
-        )
+        original = Score(name="round_trip_scorer", score=0.95, metadata={"info": "test"})
 
         # Serialize to dict
         as_dict = original.as_dict()