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

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()

braintrust/test_span_components.py

@@ -6,7 +6,6 @@ Tests serialization, deserialization, OTEL compatibility, and backward compatibi
 from uuid import uuid4
 
 import pytest
-
 from braintrust.id_gen import OTELIDGenerator
 from braintrust.span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3
 from braintrust.span_identifier_v4 import SpanComponentsV4
@@ -22,7 +21,7 @@ class TestSpanComponentsV3:
             object_id=str(uuid4()),
             row_id=str(uuid4()),
             span_id=str(uuid4()),
-            root_span_id=str(uuid4())
+            root_span_id=str(uuid4()),
         )
 
         exported = components.to_str()
@@ -39,7 +38,7 @@ class TestSpanComponentsV3:
         components = SpanComponentsV3(
             object_type=SpanObjectTypeV3.EXPERIMENT,
             object_id=str(uuid4()),
-            propagated_event={"key": "value", "nested": {"a": 1}}
+            propagated_event={"key": "value", "nested": {"a": 1}},
         )
 
         exported = components.to_str()
@@ -53,15 +52,15 @@ class TestSpanComponentsV3:
         """Test that V3 fails to preserve OTEL hex strings for 16-byte IDs (converts to UUID format)."""
         otel_gen = OTELIDGenerator()
         trace_id = otel_gen.get_trace_id()  # 32-char hex (16 bytes)
-        span_id = otel_gen.get_span_id()    # 16-char hex (8 bytes)
+        span_id = otel_gen.get_span_id()  # 16-char hex (8 bytes)
 
         # Use 16-byte hex strings for object_id and root_span_id to see UUID conversion
         components = SpanComponentsV3(
             object_type=SpanObjectTypeV3.PROJECT_LOGS,
             object_id=trace_id,  # 16-byte hex should get converted to UUID format
-            row_id='test-row-id',
-            span_id=span_id,     # 8-byte hex might be preserved
-            root_span_id=trace_id  # 16-byte hex should get converted to UUID format
+            row_id="test-row-id",
+            span_id=span_id,  # 8-byte hex might be preserved
+            root_span_id=trace_id,  # 16-byte hex should get converted to UUID format
         )
 
         exported = components.to_str()
@@ -79,14 +78,14 @@ class TestSpanComponentsV4:
         """Test that V4 preserves OTEL hex strings exactly."""
         otel_gen = OTELIDGenerator()
         trace_id = otel_gen.get_trace_id()  # 32-char hex
-        span_id = otel_gen.get_span_id()    # 16-char hex
+        span_id = otel_gen.get_span_id()  # 16-char hex
 
         components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.PROJECT_LOGS,
-            object_id='test-project-id',
-            row_id='test-row-id',
+            object_id="test-project-id",
+            row_id="test-row-id",
             span_id=span_id,
-            root_span_id=trace_id
+            root_span_id=trace_id,
         )
 
         exported = components.to_str()
@@ -108,9 +107,9 @@ class TestSpanComponentsV4:
         components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.PROJECT_LOGS,
             object_id=uuid_object_id,
-            row_id='test-row-id',
+            row_id="test-row-id",
             span_id=uuid_span_id,
-            root_span_id=uuid_root_span_id
+            root_span_id=uuid_root_span_id,
         )
 
         exported = components.to_str()
@@ -133,9 +132,9 @@ class TestSpanComponentsV4:
         components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.EXPERIMENT,
             object_id=uuid_object_id,
-            row_id='test-row-id',
+            row_id="test-row-id",
             span_id=hex_span_id,
-            root_span_id=hex_trace_id
+            root_span_id=hex_trace_id,
         )
 
         exported = components.to_str()
@@ -162,10 +161,10 @@ class TestSpanComponentsV4:
         # Create equivalent Python object
         py_components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.EXPERIMENT,
-            object_id='js-test-experiment-id',
-            row_id='js-test-row-id',
-            span_id='abcdef1234567890',
-            root_span_id='fedcba0987654321fedcba0987654321'
+            object_id="js-test-experiment-id",
+            row_id="js-test-row-id",
+            span_id="abcdef1234567890",
+            root_span_id="fedcba0987654321fedcba0987654321",
         )
 
         # Python should generate the same slug
@@ -184,8 +183,8 @@ class TestSpanComponentsV4:
         """Test V4 with additional metadata."""
         components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.PLAYGROUND_LOGS,
-            object_id='test-session-id',
-            propagated_event={"user": "test", "data": [1, 2, 3]}
+            object_id="test-session-id",
+            propagated_event={"user": "test", "data": [1, 2, 3]},
         )
 
         exported = components.to_str()
@@ -199,14 +198,14 @@ class TestSpanComponentsV4:
         """Test that non-UUID/hex strings are stored in JSON portion."""
         components = SpanComponentsV4(
             object_type=SpanObjectTypeV3.PROJECT_LOGS,
-            object_id='not-a-uuid-or-hex',  # Will be stored in JSON
+            object_id="not-a-uuid-or-hex",  # Will be stored in JSON
             # Don't test row_id alone - if present, span_id and root_span_id must also be present
         )
 
         exported = components.to_str()
         imported = SpanComponentsV4.from_str(exported)
 
-        assert imported.object_id == 'not-a-uuid-or-hex'
+        assert imported.object_id == "not-a-uuid-or-hex"
 
 
 class TestBackwardCompatibility:
@@ -221,7 +220,7 @@ class TestBackwardCompatibility:
             row_id=str(uuid4()),
             span_id=str(uuid4()),
             root_span_id=str(uuid4()),
-            propagated_event={"version": "v3"}
+            propagated_event={"version": "v3"},
         )
 
         # Serialize with V3
@@ -238,7 +237,6 @@ class TestBackwardCompatibility:
         assert v4_imported.propagated_event == v3_components.propagated_event
 
 
-
 class TestErrorHandling:
     """Test error handling and edge cases."""
 
@@ -247,7 +245,7 @@ class TestErrorHandling:
         with pytest.raises(AssertionError):
             SpanComponentsV4(
                 object_type="invalid_type",  # Should be SpanObjectTypeV3 enum
-                object_id="test-id"
+                object_id="test-id",
             )
 
     def test_missing_required_fields(self):
@@ -280,10 +278,7 @@ class TestErrorHandling:
         import base64
 
         # Create valid data then corrupt it
-        components = SpanComponentsV4(
-            object_type=SpanObjectTypeV3.PROJECT_LOGS,
-            object_id="test-id"
-        )
+        components = SpanComponentsV4(object_type=SpanObjectTypeV3.PROJECT_LOGS, object_id="test-id")
         valid_exported = components.to_str()
 
         # Decode, corrupt, re-encode
@@ -302,30 +297,21 @@ class TestObjectIdFields:
 
     def test_experiment_object_id_fields(self):
         """Test object_id_fields for experiment type."""
-        components = SpanComponentsV4(
-            object_type=SpanObjectTypeV3.EXPERIMENT,
-            object_id="test-experiment-id"
-        )
+        components = SpanComponentsV4(object_type=SpanObjectTypeV3.EXPERIMENT, object_id="test-experiment-id")
 
         fields = components.object_id_fields()
         assert fields == {"experiment_id": "test-experiment-id"}
 
     def test_project_logs_object_id_fields(self):
         """Test object_id_fields for project_logs type."""
-        components = SpanComponentsV4(
-            object_type=SpanObjectTypeV3.PROJECT_LOGS,
-            object_id="test-project-id"
-        )
+        components = SpanComponentsV4(object_type=SpanObjectTypeV3.PROJECT_LOGS, object_id="test-project-id")
 
         fields = components.object_id_fields()
         assert fields == {"project_id": "test-project-id", "log_id": "g"}
 
     def test_playground_logs_object_id_fields(self):
         """Test object_id_fields for playground_logs type."""
-        components = SpanComponentsV4(
-            object_type=SpanObjectTypeV3.PLAYGROUND_LOGS,
-            object_id="test-session-id"
-        )
+        components = SpanComponentsV4(object_type=SpanObjectTypeV3.PLAYGROUND_LOGS, object_id="test-session-id")
 
         fields = components.object_id_fields()
         assert fields == {"prompt_session_id": "test-session-id", "log_id": "x"}
@@ -333,8 +319,7 @@ class TestObjectIdFields:
     def test_object_id_fields_without_object_id(self):
         """Test that object_id_fields raises error without object_id."""
         components = SpanComponentsV4(
-            object_type=SpanObjectTypeV3.PROJECT_LOGS,
-            compute_object_metadata_args={"key": "value"}
+            object_type=SpanObjectTypeV3.PROJECT_LOGS, compute_object_metadata_args={"key": "value"}
         )
 
         with pytest.raises(Exception) as exc_info:

braintrust/util.py

@@ -2,8 +2,9 @@ import inspect
 import sys
 import threading
 import urllib.parse
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, Generic, Literal, Mapping, Optional, Set, Tuple, TypedDict, TypeVar, Union
+from typing import Any, Generic, Literal, TypedDict, TypeVar, Union
 
 from requests import HTTPError, Response
 
@@ -29,8 +30,8 @@ def coalesce(*args):
 
 
 def merge_dicts_with_paths(
-    merge_into: Dict[str, Any], merge_from: Mapping[str, Any], path: Tuple[str, ...], merge_paths: Set[Tuple[str]]
-) -> Dict[str, Any]:
+    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."""
 
@@ -50,7 +51,7 @@ def merge_dicts_with_paths(
     return merge_into
 
 
-def merge_dicts(merge_into: Dict[str, Any], merge_from: Mapping[str, Any]) -> Dict[str, Any]:
+def merge_dicts(merge_into: dict[str, Any], merge_from: Mapping[str, Any]) -> dict[str, Any]:
     """Merges merge_from into merge_into, destructively updating merge_into."""
 
     return merge_dicts_with_paths(merge_into, merge_from, (), set())
@@ -92,7 +93,7 @@ class CallerLocation(TypedDict):
     caller_lineno: int
 
 
-def get_caller_location() -> Optional[CallerLocation]:
+def get_caller_location() -> CallerLocation | None:
     frame = inspect.currentframe()
     while frame:
         frame = frame.f_back
@@ -145,7 +146,7 @@ class LazyValue(Generic[T]):
         return self._state.has_succeeded
 
     @property
-    def value(self) -> Optional[T]:
+    def value(self) -> T | None:
         return self._state.value if self._state.has_succeeded == True else None
 
     def get(self) -> T:
@@ -167,7 +168,7 @@ class LazyValue(Generic[T]):
             if self.mutex:
                 self.mutex.release()
 
-    def get_sync(self) -> Tuple[bool, Optional[T]]:
+    def get_sync(self) -> tuple[bool, T | None]:
         """Returns a tuple of (has_succeeded, value) without triggering evaluation."""
         if self._state.has_succeeded:
             # should be fine without the mutex check
@@ -206,7 +207,7 @@ def bt_iscoroutinefunction(f):
     return inspect.iscoroutinefunction(f) or inspect.isasyncgenfunction(f) or getattr(f, BT_IS_ASYNC_ATTRIBUTE, False)
 
 
-def add_azure_blob_headers(headers: Dict[str, str], url: str) -> None:
+def add_azure_blob_headers(headers: dict[str, str], url: str) -> None:
     # According to https://stackoverflow.com/questions/37824136/put-on-sas-blob-url-without-specifying-x-ms-blob-type-header,
     # there is no way to avoid including this.
     if "blob.core.windows.net" in url:

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.3.15"
+VERSION = "0.4.1"
 
 # this will be templated during the build
-GIT_COMMIT = "dcd4f5a4be171b1cac28a5eb3534e4b55420cc06"
+GIT_COMMIT = "d9c624ea93ca6bf62c2412abce1b3a2ef1a2be67"

braintrust/wrappers/_anthropic_utils.py

@@ -1,6 +1,6 @@
 """Shared utilities for Anthropic API wrappers."""
 
-from typing import Any, Dict
+from typing import Any
 
 
 class Wrapper:
@@ -13,7 +13,7 @@ class Wrapper:
         return getattr(self.__wrapped, name)
 
 
-def extract_anthropic_usage(usage: Any) -> Dict[str, float]:
+def extract_anthropic_usage(usage: Any) -> dict[str, float]:
     """Extract and normalize usage metrics from Anthropic usage object or dict.
 
     Converts Anthropic's usage format to Braintrust's standard token metric names.
@@ -29,7 +29,7 @@ def extract_anthropic_usage(usage: Any) -> Dict[str, float]:
         - prompt_cached_tokens (from cache_read_input_tokens)
         - prompt_cache_creation_tokens (from cache_creation_input_tokens)
     """
-    metrics: Dict[str, float] = {}
+    metrics: dict[str, float] = {}
 
     if not usage:
         return metrics
@@ -73,7 +73,7 @@ def extract_anthropic_usage(usage: Any) -> Dict[str, float]:
     return metrics
 
 
-def finalize_anthropic_tokens(metrics: Dict[str, float]) -> Dict[str, float]:
+def finalize_anthropic_tokens(metrics: dict[str, float]) -> dict[str, float]:
     """Finalize Anthropic token calculations.
 
     Anthropic doesn't include cache tokens in the total, so we need to sum them.

braintrust/wrappers/agno/__init__.py

@@ -21,7 +21,6 @@ Usage:
 __all__ = ["setup_agno", "wrap_agent", "wrap_function_call", "wrap_model", "wrap_team"]
 
 import logging
-from typing import Optional
 
 from braintrust.logger import NOOP_SPAN, current_span, init_logger
 
@@ -34,9 +33,9 @@ logger = logging.getLogger(__name__)
 
 
 def setup_agno(
-    api_key: Optional[str] = None,
-    project_id: Optional[str] = None,
-    project_name: Optional[str] = None,
+    api_key: str | None = None,
+    project_id: str | None = None,
+    project_name: str | None = None,
 ) -> bool:
     """
     Setup Braintrust integration with Agno. Will automatically patch Agno agents, models, and function calls for tracing.

braintrust/wrappers/agno/agent.py

@@ -1,10 +1,9 @@
 import time
 from typing import Any
 
-from wrapt import wrap_function_wrapper
-
 from braintrust.logger import start_span
 from braintrust.span_types import SpanTypeAttribute
+from wrapt import wrap_function_wrapper
 
 from .utils import (
     _aggregate_agent_chunks,

braintrust/wrappers/agno/function_call.py

@@ -1,9 +1,8 @@
 from typing import Any
 
-from wrapt import wrap_function_wrapper
-
 from braintrust.logger import start_span
 from braintrust.span_types import SpanTypeAttribute
+from wrapt import wrap_function_wrapper
 
 from .utils import is_patched
 

braintrust/wrappers/agno/model.py

@@ -5,10 +5,9 @@ ModelWrapper class for Braintrust-Agno model observability.
 import time
 from typing import Any
 
-from wrapt import wrap_function_wrapper
-
 from braintrust.logger import start_span
 from braintrust.span_types import SpanTypeAttribute
+from wrapt import wrap_function_wrapper
 
 from .utils import (
     _aggregate_model_chunks,

braintrust/wrappers/agno/team.py

@@ -1,10 +1,9 @@
 import time
 from typing import Any
 
-from wrapt import wrap_function_wrapper
-
 from braintrust.logger import start_span
 from braintrust.span_types import SpanTypeAttribute
+from wrapt import wrap_function_wrapper
 
 from .utils import (
     _aggregate_agent_chunks,