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

braintrust/otel/test_otel_bt_integration.py

@@ -8,7 +8,6 @@ when created in mixed contexts.
 import os
 
 import pytest
-
 from braintrust import current_span
 from braintrust.logger import _internal_with_memory_background_logger
 from braintrust.otel import BraintrustSpanProcessor
@@ -20,6 +19,7 @@ try:
     from opentelemetry.sdk.trace.export import SimpleSpanProcessor
     from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 except ImportError:
+
     class InMemorySpanExporter:
         def __init__(self):
             pass
@@ -44,17 +44,17 @@ class OtelFixture:
 
 @pytest.fixture
 def otel_fixture():
-    """ otel fixture configures everything we need to run mixed otel/bt tracing tests
-        that export to memory.
+    """otel fixture configures everything we need to run mixed otel/bt tracing tests
+    that export to memory.
     """
     if not OTEL_AVAILABLE:
         pytest.skip("OpenTelemetry not installed")
 
-    with preserve_env_vars('BRAINTRUST_OTEL_COMPAT', 'BRAINTRUST_API_KEY'):
+    with preserve_env_vars("BRAINTRUST_OTEL_COMPAT", "BRAINTRUST_API_KEY"):
         # 1. Set environment variable first
-        os.environ['BRAINTRUST_OTEL_COMPAT'] = 'true'
+        os.environ["BRAINTRUST_OTEL_COMPAT"] = "true"
         # Set dummy API key for tests
-        os.environ['BRAINTRUST_API_KEY'] = 'test-api-key-for-fixture'
+        os.environ["BRAINTRUST_API_KEY"] = "test-api-key-for-fixture"
 
         # 2. Set up memory logger with proper context manager
         with _internal_with_memory_background_logger() as memory_logger:
@@ -109,8 +109,8 @@ def test_mixed_otel_bt_tracing_with_bt_logger_first(otel_fixture):
     s2 = otel_spans_by_name["2"]
 
     # Verify unified trace IDs - convert OTEL trace to hex string for comparison
-    s2_trace_id = format(s2.context.trace_id, '032x')
-    s2_span_id = format(s2.context.span_id, '016x')
+    s2_trace_id = format(s2.context.trace_id, "032x")
+    s2_span_id = format(s2.context.span_id, "016x")
 
     assert s1["root_span_id"] == s2_trace_id
     assert s1["root_span_id"] == s3["root_span_id"]
@@ -151,8 +151,8 @@ def test_mixed_otel_bt_tracing_with_experiment_parent(otel_fixture):
     s1, s2, s3 = spans_by_name["1"], spans_by_name["2"], spans_by_name["3"]
 
     # Verify unified trace IDs - convert OTEL trace to hex string for comparison
-    s2_trace_id = format(s2.context.trace_id, '032x')
-    s2_span_id = format(s2.context.span_id, '016x')
+    s2_trace_id = format(s2.context.trace_id, "032x")
+    s2_span_id = format(s2.context.span_id, "016x")
 
     assert s1["root_span_id"] == s2_trace_id
     assert s1["root_span_id"] == s3["root_span_id"]
@@ -193,10 +193,10 @@ def test_mixed_otel_bt_tracing_with_otel_first(otel_fixture):
     s1, s2, s3 = spans_by_name["1"], spans_by_name["2"], spans_by_name["3"]
 
     # Verify unified trace IDs - convert OTEL traces to hex string for comparison
-    s1_trace_id = format(s1.context.trace_id, '032x')
-    s1_span_id = format(s1.context.span_id, '016x')
-    s3_trace_id = format(s3.context.trace_id, '032x')
-    s3_span_id = format(s3.context.span_id, '016x')
+    s1_trace_id = format(s1.context.trace_id, "032x")
+    s1_span_id = format(s1.context.span_id, "016x")
+    s3_trace_id = format(s3.context.trace_id, "032x")
+    s3_span_id = format(s3.context.span_id, "016x")
 
     assert s1_trace_id == s2["root_span_id"]
     assert s1_trace_id == s3_trace_id
@@ -222,14 +222,14 @@ def test_separate_traces_should_not_be_unified(otel_fixture):
     # Second trace: OTEL only
     trace2_spans = []
     with tracer.start_as_current_span("otel_trace2") as otel_span2:
-        trace2_id = format(otel_span2.context.trace_id, '032x')
+        trace2_id = format(otel_span2.context.trace_id, "032x")
         trace2_spans.append(trace2_id)
         otel_span2.set_attribute("test", "second_trace")
 
     # Third trace: OTEL root with BT child
     trace3_spans = []
     with tracer.start_as_current_span("otel_trace3_root") as otel_span3:
-        otel3_trace_id = format(otel_span3.context.trace_id, '032x')
+        otel3_trace_id = format(otel_span3.context.trace_id, "032x")
         trace3_spans.append(otel3_trace_id)
 
         # BT span inside OTEL - should inherit OTEL trace ID, not previous BT trace
@@ -286,28 +286,28 @@ def test_otel_spans_inherit_parent_attribute(otel_fixture):
         assert len(bt_spans) == 1
 
 
-
 def test_uses_braintrust_context_manager_when_otel_disabled():
     """Test that BraintrustContextManager is used when OTEL is not enabled."""
     # Ensure OTEL is disabled
-    os.environ.pop('BRAINTRUST_OTEL_COMPAT', None)
+    os.environ.pop("BRAINTRUST_OTEL_COMPAT", None)
 
     try:
         from braintrust.context import get_context_manager
+
         cm = get_context_manager()
 
         # Should be BraintrustContextManager, not OTEL ContextManager
         assert type(cm).__name__ == "BraintrustContextManager"
 
         # Verify it has the expected interface
-        assert hasattr(cm, 'get_current_span_info')
-        assert hasattr(cm, 'get_parent_span_ids')
-        assert hasattr(cm, 'set_current_span')
-        assert hasattr(cm, 'unset_current_span')
+        assert hasattr(cm, "get_current_span_info")
+        assert hasattr(cm, "get_parent_span_ids")
+        assert hasattr(cm, "set_current_span")
+        assert hasattr(cm, "unset_current_span")
 
     finally:
         # Clean up - remove any environment variable we might have set
-        os.environ.pop('BRAINTRUST_OTEL_COMPAT', None)
+        os.environ.pop("BRAINTRUST_OTEL_COMPAT", None)
 
 
 def test_uses_otel_context_manager_when_enabled():
@@ -315,21 +315,22 @@ def test_uses_otel_context_manager_when_enabled():
     if not OTEL_AVAILABLE:
         pytest.skip("OpenTelemetry not installed")
 
-    with preserve_env_vars('BRAINTRUST_OTEL_COMPAT'):
+    with preserve_env_vars("BRAINTRUST_OTEL_COMPAT"):
         # Enable OTEL
-        os.environ['BRAINTRUST_OTEL_COMPAT'] = 'true'
+        os.environ["BRAINTRUST_OTEL_COMPAT"] = "true"
 
         from braintrust.context import get_context_manager
+
         cm = get_context_manager()
 
         # Should be OTEL ContextManager, not BraintrustContextManager
         assert type(cm).__name__ == "ContextManager"
 
         # Verify it has the expected interface
-        assert hasattr(cm, 'get_current_span_info')
-        assert hasattr(cm, 'get_parent_span_ids')
-        assert hasattr(cm, 'set_current_span')
-        assert hasattr(cm, 'unset_current_span')
+        assert hasattr(cm, "get_current_span_info")
+        assert hasattr(cm, "get_parent_span_ids")
+        assert hasattr(cm, "set_current_span")
+        assert hasattr(cm, "unset_current_span")
 
 
 def test_bt_span_without_explicit_parent_inherits_from_otel(otel_fixture):
@@ -359,8 +360,8 @@ def test_bt_span_without_explicit_parent_inherits_from_otel(otel_fixture):
     otel_parent = otel_spans[0]
 
     # Convert OTEL IDs to hex for comparison
-    otel_trace_id = format(otel_parent.context.trace_id, '032x')
-    otel_span_id = format(otel_parent.context.span_id, '016x')
+    otel_trace_id = format(otel_parent.context.trace_id, "032x")
+    otel_span_id = format(otel_parent.context.span_id, "016x")
 
     # BT span should have inherited OTEL parent's trace ID as root_span_id
     assert bt_child["root_span_id"] == otel_trace_id

braintrust/parameters.py

@@ -1,6 +1,6 @@
 """Evaluation parameters support for Python SDK."""
 
-from typing import Any, Dict, Optional, TypedDict, Union
+from typing import Any, TypedDict
 
 from typing_extensions import NotRequired
 
@@ -12,17 +12,17 @@ class PromptParameter(TypedDict):
     """A prompt parameter specification."""
 
     type: str  # Literal["prompt"] but using str for flexibility
-    default: NotRequired[Optional[PromptData]]
-    description: NotRequired[Optional[str]]
+    default: NotRequired[PromptData | None]
+    description: NotRequired[str | None]
 
 
 # EvalParameters is a dict where values can be either:
 # - A PromptParameter (dict with type="prompt")
 # - A pydantic model class (typed as Any for now)
-EvalParameters = Dict[str, Union[PromptParameter, Any]]
+EvalParameters = dict[str, PromptParameter | Any]
 
 
-def _pydantic_to_json_schema(model: Any) -> Dict[str, Any]:
+def _pydantic_to_json_schema(model: Any) -> dict[str, Any]:
     """Convert a pydantic model to JSON schema."""
     if hasattr(model, "model_json_schema"):
         # pydantic 2
@@ -35,9 +35,9 @@ def _pydantic_to_json_schema(model: Any) -> Dict[str, Any]:
 
 
 def validate_parameters(
-    parameters: Dict[str, Any],
+    parameters: dict[str, Any],
     parameter_schema: EvalParameters,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Validate parameters against the schema.
 
@@ -120,7 +120,7 @@ def validate_parameters(
     return result
 
 
-def parameters_to_json_schema(parameters: EvalParameters) -> Dict[str, Any]:
+def parameters_to_json_schema(parameters: EvalParameters) -> dict[str, Any]:
     """
     Convert EvalParameters to JSON schema format for serialization.
 

braintrust/prompt.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Literal, Optional, Union
+from typing import Literal, Union
 
 from .generated_types import PromptOptions
 from .serializable_data_class import SerializableDataClass
@@ -46,17 +46,17 @@ class ImagePart(SerializableDataClass):
 
 @dataclass
 class PromptMessage(SerializableDataClass):
-    content: Union[str, List[Union[TextPart, ImagePart]]]
+    content: str | list[TextPart | ImagePart]
     role: Literal["system", "user", "assistant", "function", "tool", "model"]
-    name: Optional[str] = None
-    function_call: Union[str, FunctionCall, None] = None
-    tool_calls: Optional[List[ToolCall]] = None
+    name: str | None = None
+    function_call: str | FunctionCall | None = None
+    tool_calls: list[ToolCall] | None = None
 
 
 @dataclass
 class PromptChatBlock(SerializableDataClass):
-    messages: List[PromptMessage]
-    tools: Optional[str] = None
+    messages: list[PromptMessage]
+    tools: str | None = None
     type: Literal["chat"] = "chat"
 
 
@@ -65,20 +65,20 @@ PromptBlockData = Union[PromptCompletionBlock, PromptChatBlock]
 
 @dataclass
 class PromptData(SerializableDataClass):
-    prompt: Optional[PromptBlockData] = None
-    options: Optional[PromptOptions] = None
+    prompt: PromptBlockData | None = None
+    options: PromptOptions | None = None
 
 
 @dataclass
 class PromptSchema(SerializableDataClass):
-    id: Optional[str]
-    project_id: Optional[str]
-    _xact_id: Optional[str]
+    id: str | None
+    project_id: str | None
+    _xact_id: str | None
     name: str
     slug: str
-    description: Optional[str]
+    description: str | None
     prompt_data: PromptData
-    tags: Optional[List[str]]
+    tags: list[str] | None
 
 
 BRAINTRUST_PARAMS = ["use_cache"]

braintrust/prompt_cache/disk_cache.py

@@ -12,7 +12,8 @@ import hashlib
 import json
 import logging
 import os
-from typing import Any, Callable, Generic, Optional, TypeVar
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
 
 T = TypeVar("T")
 
@@ -36,9 +37,9 @@ class DiskCache(Generic[T]):
     def __init__(
         self,
         cache_dir: str,
-        max_size: Optional[int] = None,
-        serializer: Optional[Callable[[T], Any]] = None,
-        deserializer: Optional[Callable[[Any], T]] = None,
+        max_size: int | None = None,
+        serializer: Callable[[T], Any] | None = None,
+        deserializer: Callable[[Any], T] | None = None,
         log_warnings: bool = True,
         mkdirs: bool = True,
     ):

braintrust/prompt_cache/lru_cache.py

@@ -7,7 +7,8 @@ used items when it reaches a maximum size. The implementation uses an OrderedDic
 for O(1) access and update operations.
 """
 
-from typing import Generic, Optional, OrderedDict, TypeVar
+from collections import OrderedDict
+from typing import Generic, TypeVar
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -28,7 +29,7 @@ class LRUCache(Generic[K, V]):
                  If not specified, the cache will grow unbounded.
     """
 
-    def __init__(self, max_size: Optional[int] = None):
+    def __init__(self, max_size: int | None = None):
         self._cache: OrderedDict[K, V] = OrderedDict()
         self._max_size = max_size
 

braintrust/prompt_cache/prompt_cache.py

@@ -9,18 +9,17 @@ This allows for efficient prompt retrieval while maintaining persistence across
 The cache is keyed by project identifier (ID or name), prompt slug, and version.
 """
 
-from typing import Optional
 
 from braintrust import prompt
 from braintrust.prompt_cache import disk_cache, lru_cache
 
 
 def _create_cache_key(
-    project_id: Optional[str],
-    project_name: Optional[str],
-    slug: Optional[str],
+    project_id: str | None,
+    project_name: str | None,
+    slug: str | None,
     version: str = "latest",
-    id: Optional[str] = None,
+    id: str | None = None,
 ) -> str:
     """Creates a unique cache key from project identifier, slug and version, or from ID."""
     if id:
@@ -47,7 +46,7 @@ class PromptCache:
     def __init__(
         self,
         memory_cache: lru_cache.LRUCache[str, prompt.PromptSchema],
-        disk_cache: Optional[disk_cache.DiskCache[prompt.PromptSchema]] = None,
+        disk_cache: disk_cache.DiskCache[prompt.PromptSchema] | None = None,
     ):
         """
         Initialize the prompt cache.
@@ -61,11 +60,11 @@ class PromptCache:
 
     def get(
         self,
-        slug: Optional[str] = None,
+        slug: str | None = None,
         version: str = "latest",
-        project_id: Optional[str] = None,
-        project_name: Optional[str] = None,
-        id: Optional[str] = None,
+        project_id: str | None = None,
+        project_name: str | None = None,
+        id: str | None = None,
     ) -> prompt.PromptSchema:
         """
         Retrieve a prompt from the cache.
@@ -107,11 +106,11 @@ class PromptCache:
     def set(
         self,
         value: prompt.PromptSchema,
-        slug: Optional[str] = None,
+        slug: str | None = None,
         version: str = "latest",
-        project_id: Optional[str] = None,
-        project_name: Optional[str] = None,
-        id: Optional[str] = None,
+        project_id: str | None = None,
+        project_name: str | None = None,
+        id: str | None = None,
     ) -> None:
         """
         Store a prompt in the cache.

braintrust/queue.py

@@ -1,6 +1,6 @@
 import threading
 from collections import deque
-from typing import List, Optional, TypeVar
+from typing import TypeVar
 
 from .util import eprint
 
@@ -46,7 +46,7 @@ class LogQueue:
         with self._mutex:
             self._enforce_size_limit = enforce
 
-    def put(self, item: T) -> List[T]:
+    def put(self, item: T) -> list[T]:
         """
         Put an item in the queue.
 
@@ -76,7 +76,7 @@ class LogQueue:
 
         return dropped
 
-    def drain_all(self) -> List[T]:
+    def drain_all(self) -> list[T]:
         """
         Drain all items from the queue.
 
@@ -105,7 +105,7 @@ class LogQueue:
         """
         return len(self._queue)
 
-    def wait_for_items(self, timeout: Optional[float] = None) -> bool:
+    def wait_for_items(self, timeout: float | None = None) -> bool:
         """
         Will block until the queue has at least one item in it. Might be empty by the time
         you read though.

braintrust/score.py

@@ -2,7 +2,7 @@ import dataclasses
 import inspect
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional
+from typing import Any
 
 from .serializable_data_class import SerializableDataClass
 
@@ -22,14 +22,14 @@ class Score(SerializableDataClass):
     name: str
     """The name of the score. This should be a unique name for the scorer."""
 
-    score: Optional[float] = None
+    score: float | None = None
     """The score for the evaluation. This should be a float between 0 and 1. If the score is None, the evaluation is considered to be skipped."""
 
-    metadata: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
     """Metadata for the score. This can be used to store additional information about the score."""
 
     # DEPRECATION_NOTICE: this field is deprecated, as errors are propagated up to the caller.
-    error: Optional[Exception] = None
+    error: Exception | None = None
     """Deprecated: The error field is deprecated, as errors are now propagated to the caller. The field will be removed in a future version of the library."""
 
     def as_dict(self):

braintrust/serializable_data_class.py

@@ -1,6 +1,6 @@
 import dataclasses
 import json
-from typing import Dict, Union, get_origin
+from typing import Union, get_origin
 
 
 class SerializableDataClass:
@@ -16,15 +16,15 @@ class SerializableDataClass:
         return getattr(self, item)
 
     @classmethod
-    def from_dict(cls, d: Dict):
+    def from_dict(cls, d: dict):
         """Deserialize the object from a dictionary. This method
         is shallow and will not call from_dict() on nested objects."""
-        fields = set(f.name for f in dataclasses.fields(cls))
+        fields = {f.name for f in dataclasses.fields(cls)}
         filtered = {k: v for k, v in d.items() if k in fields}
         return cls(**filtered)
 
     @classmethod
-    def from_dict_deep(cls, d: Dict):
+    def from_dict_deep(cls, d: dict):
         """Deserialize the object from a dictionary. This method
         is deep and will call from_dict_deep() on nested objects."""
         fields = {f.name: f for f in dataclasses.fields(cls)}

braintrust/span_identifier_v1.py

@@ -5,7 +5,6 @@
 import base64
 import dataclasses
 from enum import Enum, auto
-from typing import Optional
 from uuid import UUID
 
 
@@ -50,7 +49,7 @@ class SpanRowIdsV1:
 class SpanComponentsV1:
     object_type: SpanObjectTypeV1
     object_id: str
-    row_ids: Optional[SpanRowIdsV1] = None
+    row_ids: SpanRowIdsV1 | None = None
 
     def __post_init__(self):
         assert isinstance(self.object_type, SpanObjectTypeV1)

braintrust/span_identifier_v2.py

@@ -6,7 +6,6 @@ import base64
 import dataclasses
 import json
 from enum import Enum
-from typing import Dict, Optional
 from uuid import UUID
 
 from .span_identifier_v1 import SpanComponentsV1
@@ -54,9 +53,9 @@ class SpanRowIdsV2:
 @dataclasses.dataclass
 class SpanComponentsV2:
     object_type: SpanObjectTypeV2
-    object_id: Optional[str] = None
-    compute_object_metadata_args: Optional[Dict] = None
-    row_ids: Optional[SpanRowIdsV2] = None
+    object_id: str | None = None
+    compute_object_metadata_args: dict | None = None
+    row_ids: SpanRowIdsV2 | None = None
 
     def __post_init__(self):
         assert isinstance(self.object_type, SpanObjectTypeV2)

braintrust/span_identifier_v3.py

@@ -6,7 +6,6 @@ import base64
 import dataclasses
 import json
 from enum import Enum
-from typing import Dict, Optional, Union
 from uuid import UUID
 
 from .span_identifier_v2 import SpanComponentsV2
@@ -59,16 +58,16 @@ class SpanComponentsV3:
     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):
         assert isinstance(self.object_type, SpanObjectTypeV3)
@@ -173,7 +172,7 @@ class SpanComponentsV3:
         except Exception:
             raise Exception(INVALID_ENCODING_ERRMSG)
 
-    def object_id_fields(self) -> Dict[str, str]:
+    def object_id_fields(self) -> dict[str, str]:
         if not self.object_id:
             raise Exception(
                 "Impossible: cannot invoke `object_id_fields` unless SpanComponentsV3 is initialized with an `object_id`"
@@ -192,7 +191,7 @@ class SpanComponentsV3:
         return self.to_str()
 
     @staticmethod
-    def _from_json_obj(json_obj: Dict) -> "SpanComponentsV3":
+    def _from_json_obj(json_obj: dict) -> "SpanComponentsV3":
         kwargs = {
             **json_obj,
             "object_type": SpanObjectTypeV3(json_obj["object_type"]),
@@ -200,7 +199,7 @@ class SpanComponentsV3:
         return SpanComponentsV3(**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.
 
@@ -235,17 +234,21 @@ def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
         # Handle row_ids if present
         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,
+                }
+            )
 
         # Include propagated_event if present
         if "propagated_event" in parent: