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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (83) hide show
  1. braintrust/__init__.py +4 -0
  2. braintrust/_generated_types.py +1200 -611
  3. braintrust/audit.py +2 -2
  4. braintrust/cli/eval.py +6 -7
  5. braintrust/cli/push.py +11 -11
  6. braintrust/conftest.py +1 -0
  7. braintrust/context.py +12 -17
  8. braintrust/contrib/temporal/__init__.py +16 -27
  9. braintrust/contrib/temporal/test_temporal.py +8 -3
  10. braintrust/devserver/auth.py +8 -8
  11. braintrust/devserver/cache.py +3 -4
  12. braintrust/devserver/cors.py +8 -7
  13. braintrust/devserver/dataset.py +3 -5
  14. braintrust/devserver/eval_hooks.py +7 -6
  15. braintrust/devserver/schemas.py +22 -19
  16. braintrust/devserver/server.py +19 -12
  17. braintrust/devserver/test_cached_login.py +4 -4
  18. braintrust/framework.py +128 -140
  19. braintrust/framework2.py +88 -87
  20. braintrust/functions/invoke.py +93 -53
  21. braintrust/functions/stream.py +3 -2
  22. braintrust/generated_types.py +17 -1
  23. braintrust/git_fields.py +11 -11
  24. braintrust/gitutil.py +2 -3
  25. braintrust/graph_util.py +10 -10
  26. braintrust/id_gen.py +2 -2
  27. braintrust/logger.py +346 -357
  28. braintrust/merge_row_batch.py +10 -9
  29. braintrust/oai.py +107 -24
  30. braintrust/otel/__init__.py +49 -49
  31. braintrust/otel/context.py +16 -30
  32. braintrust/otel/test_distributed_tracing.py +14 -11
  33. braintrust/otel/test_otel_bt_integration.py +32 -31
  34. braintrust/parameters.py +8 -8
  35. braintrust/prompt.py +14 -14
  36. braintrust/prompt_cache/disk_cache.py +5 -4
  37. braintrust/prompt_cache/lru_cache.py +3 -2
  38. braintrust/prompt_cache/prompt_cache.py +13 -14
  39. braintrust/queue.py +4 -4
  40. braintrust/score.py +4 -4
  41. braintrust/serializable_data_class.py +4 -4
  42. braintrust/span_identifier_v1.py +1 -2
  43. braintrust/span_identifier_v2.py +3 -4
  44. braintrust/span_identifier_v3.py +23 -20
  45. braintrust/span_identifier_v4.py +34 -25
  46. braintrust/test_framework.py +16 -6
  47. braintrust/test_helpers.py +5 -5
  48. braintrust/test_id_gen.py +2 -3
  49. braintrust/test_otel.py +61 -53
  50. braintrust/test_queue.py +0 -1
  51. braintrust/test_score.py +1 -3
  52. braintrust/test_span_components.py +29 -44
  53. braintrust/util.py +9 -8
  54. braintrust/version.py +2 -2
  55. braintrust/wrappers/_anthropic_utils.py +4 -4
  56. braintrust/wrappers/agno/__init__.py +3 -4
  57. braintrust/wrappers/agno/agent.py +1 -2
  58. braintrust/wrappers/agno/function_call.py +1 -2
  59. braintrust/wrappers/agno/model.py +1 -2
  60. braintrust/wrappers/agno/team.py +1 -2
  61. braintrust/wrappers/agno/utils.py +12 -12
  62. braintrust/wrappers/anthropic.py +7 -8
  63. braintrust/wrappers/claude_agent_sdk/__init__.py +3 -4
  64. braintrust/wrappers/claude_agent_sdk/_wrapper.py +29 -27
  65. braintrust/wrappers/dspy.py +15 -17
  66. braintrust/wrappers/google_genai/__init__.py +16 -16
  67. braintrust/wrappers/langchain.py +22 -24
  68. braintrust/wrappers/litellm.py +4 -3
  69. braintrust/wrappers/openai.py +15 -15
  70. braintrust/wrappers/pydantic_ai.py +1204 -0
  71. braintrust/wrappers/test_agno.py +0 -1
  72. braintrust/wrappers/test_dspy.py +0 -1
  73. braintrust/wrappers/test_google_genai.py +2 -3
  74. braintrust/wrappers/test_litellm.py +0 -1
  75. braintrust/wrappers/test_oai_attachments.py +322 -0
  76. braintrust/wrappers/test_pydantic_ai_integration.py +1788 -0
  77. braintrust/wrappers/{test_pydantic_ai.py → test_pydantic_ai_wrap_openai.py} +1 -2
  78. {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/METADATA +3 -2
  79. braintrust-0.4.0.dist-info/RECORD +120 -0
  80. braintrust-0.3.14.dist-info/RECORD +0 -117
  81. {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/WHEEL +0 -0
  82. {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/entry_points.txt +0 -0
  83. {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/top_level.txt +0 -0
@@ -5,7 +5,6 @@
5
5
  import base64
6
6
  import dataclasses
7
7
  from enum import Enum, auto
8
- from typing import Optional
9
8
  from uuid import UUID
10
9
 
11
10
 
@@ -50,7 +49,7 @@ class SpanRowIdsV1:
50
49
  class SpanComponentsV1:
51
50
  object_type: SpanObjectTypeV1
52
51
  object_id: str
53
- row_ids: Optional[SpanRowIdsV1] = None
52
+ row_ids: SpanRowIdsV1 | None = None
54
53
 
55
54
  def __post_init__(self):
56
55
  assert isinstance(self.object_type, SpanObjectTypeV1)
@@ -6,7 +6,6 @@ import base64
6
6
  import dataclasses
7
7
  import json
8
8
  from enum import Enum
9
- from typing import Dict, Optional
10
9
  from uuid import UUID
11
10
 
12
11
  from .span_identifier_v1 import SpanComponentsV1
@@ -54,9 +53,9 @@ class SpanRowIdsV2:
54
53
  @dataclasses.dataclass
55
54
  class SpanComponentsV2:
56
55
  object_type: SpanObjectTypeV2
57
- object_id: Optional[str] = None
58
- compute_object_metadata_args: Optional[Dict] = None
59
- row_ids: Optional[SpanRowIdsV2] = None
56
+ object_id: str | None = None
57
+ compute_object_metadata_args: dict | None = None
58
+ row_ids: SpanRowIdsV2 | None = None
60
59
 
61
60
  def __post_init__(self):
62
61
  assert isinstance(self.object_type, SpanObjectTypeV2)
@@ -6,7 +6,6 @@ import base64
6
6
  import dataclasses
7
7
  import json
8
8
  from enum import Enum
9
- from typing import Dict, Optional, Union
10
9
  from uuid import UUID
11
10
 
12
11
  from .span_identifier_v2 import SpanComponentsV2
@@ -59,16 +58,16 @@ class SpanComponentsV3:
59
58
  object_type: SpanObjectTypeV3
60
59
 
61
60
  # Must provide one or the other.
62
- object_id: Optional[str] = None
63
- compute_object_metadata_args: Optional[Dict] = None
61
+ object_id: str | None = None
62
+ compute_object_metadata_args: dict | None = None
64
63
 
65
64
  # Either all of these must be provided or none.
66
- row_id: Optional[str] = None
67
- span_id: Optional[str] = None
68
- root_span_id: Optional[str] = None
65
+ row_id: str | None = None
66
+ span_id: str | None = None
67
+ root_span_id: str | None = None
69
68
 
70
69
  # Additional span properties.
71
- propagated_event: Optional[Dict] = None
70
+ propagated_event: dict | None = None
72
71
 
73
72
  def __post_init__(self):
74
73
  assert isinstance(self.object_type, SpanObjectTypeV3)
@@ -173,7 +172,7 @@ class SpanComponentsV3:
173
172
  except Exception:
174
173
  raise Exception(INVALID_ENCODING_ERRMSG)
175
174
 
176
- def object_id_fields(self) -> Dict[str, str]:
175
+ def object_id_fields(self) -> dict[str, str]:
177
176
  if not self.object_id:
178
177
  raise Exception(
179
178
  "Impossible: cannot invoke `object_id_fields` unless SpanComponentsV3 is initialized with an `object_id`"
@@ -192,7 +191,7 @@ class SpanComponentsV3:
192
191
  return self.to_str()
193
192
 
194
193
  @staticmethod
195
- def _from_json_obj(json_obj: Dict) -> "SpanComponentsV3":
194
+ def _from_json_obj(json_obj: dict) -> "SpanComponentsV3":
196
195
  kwargs = {
197
196
  **json_obj,
198
197
  "object_type": SpanObjectTypeV3(json_obj["object_type"]),
@@ -200,7 +199,7 @@ class SpanComponentsV3:
200
199
  return SpanComponentsV3(**kwargs)
201
200
 
202
201
 
203
- def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
202
+ def parse_parent(parent: str | dict | None) -> str | None:
204
203
  """
205
204
  Parse a parent object into a string representation.
206
205
 
@@ -235,17 +234,21 @@ def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
235
234
  # Handle row_ids if present
236
235
  row_ids = parent.get("row_ids")
237
236
  if row_ids:
238
- kwargs.update({
239
- "row_id": row_ids.get("id"),
240
- "span_id": row_ids.get("span_id"),
241
- "root_span_id": row_ids.get("root_span_id"),
242
- })
237
+ kwargs.update(
238
+ {
239
+ "row_id": row_ids.get("id"),
240
+ "span_id": row_ids.get("span_id"),
241
+ "root_span_id": row_ids.get("root_span_id"),
242
+ }
243
+ )
243
244
  else:
244
- kwargs.update({
245
- "row_id": None,
246
- "span_id": None,
247
- "root_span_id": None,
248
- })
245
+ kwargs.update(
246
+ {
247
+ "row_id": None,
248
+ "span_id": None,
249
+ "root_span_id": None,
250
+ }
251
+ )
249
252
 
250
253
  # Include propagated_event if present
251
254
  if "propagated_event" in parent:
@@ -5,7 +5,6 @@ import base64
5
5
  import dataclasses
6
6
  import json
7
7
  from enum import Enum
8
- from typing import Dict, Optional, Union
9
8
 
10
9
  from .span_identifier_v3 import (
11
10
  SpanComponentsV3,
@@ -14,6 +13,7 @@ from .span_identifier_v3 import (
14
13
 
15
14
  ENCODING_VERSION_NUMBER_V4 = 4
16
15
 
16
+
17
17
  def _try_make_hex_trace_id(s):
18
18
  """Try to convert hex string to 16-byte binary (for trace IDs)"""
19
19
  try:
@@ -25,6 +25,7 @@ def _try_make_hex_trace_id(s):
25
25
  pass
26
26
  return None, False
27
27
 
28
+
28
29
  def _try_make_hex_span_id(s):
29
30
  """Try to convert hex string to 8-byte binary (for span IDs)"""
30
31
  try:
@@ -36,14 +37,17 @@ def _try_make_hex_span_id(s):
36
37
  pass
37
38
  return None, False
38
39
 
40
+
39
41
  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."
40
42
 
43
+
41
44
  class Fields(Enum):
42
45
  OBJECT_ID = 1
43
46
  ROW_ID = 2
44
47
  SPAN_ID = 3 # 8-byte hex
45
48
  ROOT_SPAN_ID = 4 # 16-byte hex
46
49
 
50
+
47
51
  _FIELDS_ID_TO_NAME = {
48
52
  Fields.OBJECT_ID: "object_id",
49
53
  Fields.ROW_ID: "row_id",
@@ -57,16 +61,16 @@ class SpanComponentsV4:
57
61
  object_type: SpanObjectTypeV3
58
62
 
59
63
  # Must provide one or the other.
60
- object_id: Optional[str] = None
61
- compute_object_metadata_args: Optional[Dict] = None
64
+ object_id: str | None = None
65
+ compute_object_metadata_args: dict | None = None
62
66
 
63
67
  # Either all of these must be provided or none.
64
- row_id: Optional[str] = None
65
- span_id: Optional[str] = None
66
- root_span_id: Optional[str] = None
68
+ row_id: str | None = None
69
+ span_id: str | None = None
70
+ root_span_id: str | None = None
67
71
 
68
72
  # Additional span properties.
69
- propagated_event: Optional[Dict] = None
73
+ propagated_event: dict | None = None
70
74
 
71
75
  def __post_init__(self):
72
76
  # Reuse V3 validation logic
@@ -98,10 +102,12 @@ class SpanComponentsV4:
98
102
  )
99
103
  json_obj = {k: v for k, v in json_obj.items() if v is not None}
100
104
 
101
- raw_bytes = bytes([
102
- ENCODING_VERSION_NUMBER_V4,
103
- self.object_type.value,
104
- ])
105
+ raw_bytes = bytes(
106
+ [
107
+ ENCODING_VERSION_NUMBER_V4,
108
+ self.object_type.value,
109
+ ]
110
+ )
105
111
 
106
112
  hex_entries = []
107
113
 
@@ -199,7 +205,7 @@ class SpanComponentsV4:
199
205
  except Exception:
200
206
  raise Exception(INVALID_ENCODING_ERRMSG_V4)
201
207
 
202
- def object_id_fields(self) -> Dict[str, str]:
208
+ def object_id_fields(self) -> dict[str, str]:
203
209
  # Reuse V3 logic
204
210
  if not self.object_id:
205
211
  raise Exception(
@@ -218,7 +224,7 @@ class SpanComponentsV4:
218
224
  return self.to_str()
219
225
 
220
226
  @staticmethod
221
- def _from_json_obj(json_obj: Dict) -> "SpanComponentsV4":
227
+ def _from_json_obj(json_obj: dict) -> "SpanComponentsV4":
222
228
  kwargs = {
223
229
  **json_obj,
224
230
  "object_type": SpanObjectTypeV3(json_obj["object_type"]),
@@ -226,8 +232,7 @@ class SpanComponentsV4:
226
232
  return SpanComponentsV4(**kwargs)
227
233
 
228
234
 
229
-
230
- def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
235
+ def parse_parent(parent: str | dict | None) -> str | None:
231
236
  """Parse a parent object into a string representation using V4 format."""
232
237
  # Reuse V3 logic but with V4 components
233
238
  if isinstance(parent, str):
@@ -250,17 +255,21 @@ def parse_parent(parent: Union[str, Dict, None]) -> Optional[str]:
250
255
 
251
256
  row_ids = parent.get("row_ids")
252
257
  if row_ids:
253
- kwargs.update({
254
- "row_id": row_ids.get("id"),
255
- "span_id": row_ids.get("span_id"),
256
- "root_span_id": row_ids.get("root_span_id"),
257
- })
258
+ kwargs.update(
259
+ {
260
+ "row_id": row_ids.get("id"),
261
+ "span_id": row_ids.get("span_id"),
262
+ "root_span_id": row_ids.get("root_span_id"),
263
+ }
264
+ )
258
265
  else:
259
- kwargs.update({
260
- "row_id": None,
261
- "span_id": None,
262
- "root_span_id": None,
263
- })
266
+ kwargs.update(
267
+ {
268
+ "row_id": None,
269
+ "span_id": None,
270
+ "root_span_id": None,
271
+ }
272
+ )
264
273
 
265
274
  if "propagated_event" in parent:
266
275
  kwargs["propagated_event"] = parent.get("propagated_event")
@@ -240,12 +240,15 @@ async def test_hooks_trial_index_multiple_inputs():
240
240
  assert sorted(input_1_trials) == [0, 1]
241
241
  assert sorted(input_2_trials) == [0, 1]
242
242
 
243
+
243
244
  @pytest.fixture
244
245
  def simple_scorer():
245
246
  def simple_scorer_function(input, output, expected):
246
247
  return {"name": "simple_scorer", "score": 0.8}
248
+
247
249
  return simple_scorer_function
248
250
 
251
+
249
252
  @pytest.mark.asyncio
250
253
  async def test_eval_no_send_logs_true(with_memory_logger, simple_scorer):
251
254
  """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):
286
289
 
287
290
  @pytest.mark.asyncio
288
291
  async def test_hooks_tags_append(with_memory_logger, with_simulate_login, simple_scorer):
289
- """ Test that hooks.tags can be appended to and logged. """
292
+ """Test that hooks.tags can be appended to and logged."""
290
293
 
291
294
  initial_tags = ["cookies n cream"]
292
295
  appended_tags = ["chocolate", "vanilla", "strawberry"]
@@ -321,9 +324,12 @@ async def test_hooks_tags_append(with_memory_logger, with_simulate_login, simple
321
324
 
322
325
 
323
326
  @pytest.mark.asyncio
324
- @pytest.mark.parametrize(("tags", "expected_tags"), [(None, None),([], None), (["chocolate", "vanilla", "strawberry"], ["chocolate", "vanilla", "strawberry"])])
327
+ @pytest.mark.parametrize(
328
+ ("tags", "expected_tags"),
329
+ [(None, None), ([], None), (["chocolate", "vanilla", "strawberry"], ["chocolate", "vanilla", "strawberry"])],
330
+ )
325
331
  async def test_hooks_tags_list(with_memory_logger, with_simulate_login, simple_scorer, tags, expected_tags):
326
- """ Test that hooks.tags can be set to a list. """
332
+ """Test that hooks.tags can be set to a list."""
327
333
 
328
334
  def task_with_hooks(input, hooks):
329
335
  hooks.tags = tags
@@ -351,9 +357,10 @@ async def test_hooks_tags_list(with_memory_logger, with_simulate_login, simple_s
351
357
  assert len(root_span) == 1
352
358
  assert root_span[0].get("tags") == expected_tags
353
359
 
360
+
354
361
  @pytest.mark.asyncio
355
362
  async def test_hooks_tags_with_failing_scorer(with_memory_logger, with_simulate_login, simple_scorer):
356
- """ Test that hooks.tags can be set to a list. """
363
+ """Test that hooks.tags can be set to a list."""
357
364
 
358
365
  expected_tags = ["chocolate", "vanilla", "strawberry"]
359
366
 
@@ -386,9 +393,11 @@ async def test_hooks_tags_with_failing_scorer(with_memory_logger, with_simulate_
386
393
  assert len(root_span) == 1
387
394
  assert root_span[0].get("tags") == expected_tags
388
395
 
396
+
389
397
  @pytest.mark.asyncio
390
398
  async def test_hooks_tags_with_invalid_type(with_memory_logger, with_simulate_login, simple_scorer):
391
- """ Test that result contains an error for cases where hooks.tags is set to an invalid type. """
399
+ """Test that result contains an error for cases where hooks.tags is set to an invalid type."""
400
+
392
401
  def task_with_hooks(input, hooks):
393
402
  hooks.tags = 123
394
403
  return input
@@ -411,7 +420,8 @@ async def test_hooks_tags_with_invalid_type(with_memory_logger, with_simulate_lo
411
420
 
412
421
  @pytest.mark.asyncio
413
422
  async def test_hooks_without_setting_tags(with_memory_logger, with_simulate_login, simple_scorer):
414
- """ Test where hooks.tags is not set """
423
+ """Test where hooks.tags is not set"""
424
+
415
425
  def task_with_hooks(input, hooks):
416
426
  return input
417
427
 
@@ -2,7 +2,6 @@ import os
2
2
  from contextlib import contextmanager
3
3
 
4
4
  import pytest
5
-
6
5
  from braintrust import logger
7
6
  from braintrust.logger import ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata
8
7
  from braintrust.util import LazyValue
@@ -15,10 +14,8 @@ TEST_ORG_NAME = "test-org-name"
15
14
  def has_devserver_installed() -> bool:
16
15
  """Check if devserver dependencies (starlette, uvicorn) are installed."""
17
16
  import importlib.util
18
- return (
19
- importlib.util.find_spec("starlette") is not None
20
- and importlib.util.find_spec("uvicorn") is not None
21
- )
17
+
18
+ return importlib.util.find_spec("starlette") is not None and importlib.util.find_spec("uvicorn") is not None
22
19
 
23
20
 
24
21
  def simulate_login() -> None:
@@ -68,12 +65,14 @@ def with_memory_logger():
68
65
  # Clean up global state to prevent test contamination
69
66
  logger._state.reset_parent_state()
70
67
 
68
+
71
69
  @pytest.fixture
72
70
  def memory_logger():
73
71
  with logger._internal_with_memory_background_logger() as bgl:
74
72
  yield bgl
75
73
  logger._state.current_experiment = None
76
74
 
75
+
77
76
  @contextmanager
78
77
  def preserve_env_vars(*vars):
79
78
  original_env = {v: os.environ.get(v) for v in vars}
@@ -114,6 +113,7 @@ def init_test_logger(project_name: str):
114
113
  logger._compute_logger_metadata = fake_compute_logger_metadata
115
114
  return l
116
115
 
116
+
117
117
  def init_test_exp(experiment_name: str, project_name: str = None):
118
118
  """
119
119
  Initialize an experiment for testing with fake project and experiment metadata.
braintrust/test_id_gen.py CHANGED
@@ -1,9 +1,7 @@
1
-
2
1
  import os
3
2
  import uuid
4
3
 
5
4
  import pytest
6
-
7
5
  from braintrust import id_gen
8
6
 
9
7
 
@@ -76,7 +74,8 @@ def test_id_get_env_var(reset_id_generator_state):
76
74
 
77
75
 
78
76
  def _is_hex(s):
79
- return all(c in '0123456789abcdef' for c in s.lower())
77
+ return all(c in "0123456789abcdef" for c in s.lower())
78
+
80
79
 
81
80
  def _assert_is_hex(x):
82
81
  assert _is_hex(x)
braintrust/test_otel.py CHANGED
@@ -274,12 +274,11 @@ class TestSpanFiltering:
274
274
  except ImportError:
275
275
  pytest.skip("OpenTelemetry SDK not fully installed, skipping AISpanProcessor tests")
276
276
 
277
+ from braintrust.otel import AISpanProcessor
277
278
  from opentelemetry.sdk.trace import TracerProvider
278
279
  from opentelemetry.sdk.trace.export import SimpleSpanProcessor
279
280
  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
280
281
 
281
- from braintrust.otel import AISpanProcessor
282
-
283
282
  self.memory_exporter = InMemorySpanExporter()
284
283
  self.provider = TracerProvider()
285
284
 
@@ -403,12 +402,11 @@ class TestSpanFiltering:
403
402
  return None # Don't influence decision
404
403
 
405
404
  # Create processor with custom filter
405
+ from braintrust.otel import AISpanProcessor
406
406
  from opentelemetry.sdk.trace import TracerProvider
407
407
  from opentelemetry.sdk.trace.export import SimpleSpanProcessor
408
408
  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
409
409
 
410
- from braintrust.otel import AISpanProcessor
411
-
412
410
  memory_exporter = InMemorySpanExporter()
413
411
  processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
414
412
  provider = TracerProvider()
@@ -435,12 +433,11 @@ class TestSpanFiltering:
435
433
  return None # Don't influence decision
436
434
 
437
435
  # Create processor with custom filter
436
+ from braintrust.otel import AISpanProcessor
438
437
  from opentelemetry.sdk.trace import TracerProvider
439
438
  from opentelemetry.sdk.trace.export import SimpleSpanProcessor
440
439
  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
441
440
 
442
- from braintrust.otel import AISpanProcessor
443
-
444
441
  memory_exporter = InMemorySpanExporter()
445
442
  processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
446
443
  provider = TracerProvider()
@@ -465,12 +462,11 @@ class TestSpanFiltering:
465
462
  return None # Always defer to default logic
466
463
 
467
464
  # Create processor with custom filter
465
+ from braintrust.otel import AISpanProcessor
468
466
  from opentelemetry.sdk.trace import TracerProvider
469
467
  from opentelemetry.sdk.trace.export import SimpleSpanProcessor
470
468
  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
471
469
 
472
- from braintrust.otel import AISpanProcessor
473
-
474
470
  memory_exporter = InMemorySpanExporter()
475
471
  processor = AISpanProcessor(SimpleSpanProcessor(memory_exporter), custom_filter=custom_filter)
476
472
  provider = TracerProvider()
@@ -492,12 +488,11 @@ class TestSpanFiltering:
492
488
 
493
489
  def test_filtering_vs_unfiltered_comparison(self):
494
490
  # Set up two separate exporters and processors
491
+ from braintrust.otel import AISpanProcessor
495
492
  from opentelemetry.sdk.trace import TracerProvider
496
493
  from opentelemetry.sdk.trace.export import SimpleSpanProcessor
497
494
  from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
498
495
 
499
- from braintrust.otel import AISpanProcessor
500
-
501
496
  all_spans_exporter = InMemorySpanExporter()
502
497
  filtered_spans_exporter = InMemorySpanExporter()
503
498
 
@@ -569,49 +564,58 @@ def test_parent_from_headers_invalid_inputs():
569
564
  assert result is None
570
565
 
571
566
  # Test 2: Invalid traceparent (malformed)
572
- result = parent_from_headers({'traceparent': 'invalid'})
567
+ result = parent_from_headers({"traceparent": "invalid"})
573
568
  assert result is None
574
569
 
575
570
  # Test 3: Valid traceparent but invalid braintrust.parent format
576
- result = parent_from_headers({
577
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
578
- 'baggage': 'braintrust.parent=invalid_format'
579
- })
571
+ result = parent_from_headers(
572
+ {
573
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
574
+ "baggage": "braintrust.parent=invalid_format",
575
+ }
576
+ )
580
577
  assert result is None
581
578
 
582
579
  # Test 4: Empty project_id
583
- result = parent_from_headers({
584
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
585
- 'baggage': 'braintrust.parent=project_id:'
586
- })
580
+ result = parent_from_headers(
581
+ {
582
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
583
+ "baggage": "braintrust.parent=project_id:",
584
+ }
585
+ )
587
586
  assert result is None
588
587
 
589
588
  # Test 5: Empty project_name
590
- result = parent_from_headers({
591
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
592
- 'baggage': 'braintrust.parent=project_name:'
593
- })
589
+ result = parent_from_headers(
590
+ {
591
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
592
+ "baggage": "braintrust.parent=project_name:",
593
+ }
594
+ )
594
595
  assert result is None
595
596
 
596
597
  # Test 6: Empty experiment_id
597
- result = parent_from_headers({
598
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
599
- 'baggage': 'braintrust.parent=experiment_id:'
600
- })
598
+ result = parent_from_headers(
599
+ {
600
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
601
+ "baggage": "braintrust.parent=experiment_id:",
602
+ }
603
+ )
601
604
  assert result is None
602
605
 
603
606
  # Test 7: Invalid trace_id length (too short)
604
- result = parent_from_headers({
605
- 'traceparent': '00-4bf92f3577b34da6-00f067aa0ba902b7-01',
606
- 'baggage': 'braintrust.parent=project_name:test'
607
- })
607
+ result = parent_from_headers(
608
+ {"traceparent": "00-4bf92f3577b34da6-00f067aa0ba902b7-01", "baggage": "braintrust.parent=project_name:test"}
609
+ )
608
610
  assert result is None
609
611
 
610
612
  # Test 8: Invalid span_id length (too short)
611
- result = parent_from_headers({
612
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa-01',
613
- 'baggage': 'braintrust.parent=project_name:test'
614
- })
613
+ result = parent_from_headers(
614
+ {
615
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa-01",
616
+ "baggage": "braintrust.parent=project_name:test",
617
+ }
618
+ )
615
619
  assert result is None
616
620
 
617
621
 
@@ -623,29 +627,35 @@ def test_parent_from_headers_valid_input():
623
627
  from braintrust.otel import parent_from_headers
624
628
 
625
629
  # Test with valid project_name
626
- result = parent_from_headers({
627
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
628
- 'baggage': 'braintrust.parent=project_name:test-project'
629
- })
630
+ result = parent_from_headers(
631
+ {
632
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
633
+ "baggage": "braintrust.parent=project_name:test-project",
634
+ }
635
+ )
630
636
  assert result is not None
631
637
  # Result is base64 encoded, so just check it's a non-empty string
632
638
  assert isinstance(result, str)
633
639
  assert len(result) > 0
634
640
 
635
641
  # Test with valid project_id
636
- result = parent_from_headers({
637
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
638
- 'baggage': 'braintrust.parent=project_id:abc123'
639
- })
642
+ result = parent_from_headers(
643
+ {
644
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
645
+ "baggage": "braintrust.parent=project_id:abc123",
646
+ }
647
+ )
640
648
  assert result is not None
641
649
  assert isinstance(result, str)
642
650
  assert len(result) > 0
643
651
 
644
652
  # Test with valid experiment_id
645
- result = parent_from_headers({
646
- 'traceparent': '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01',
647
- 'baggage': 'braintrust.parent=experiment_id:exp-456'
648
- })
653
+ result = parent_from_headers(
654
+ {
655
+ "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
656
+ "baggage": "braintrust.parent=experiment_id:exp-456",
657
+ }
658
+ )
649
659
  assert result is not None
650
660
  assert isinstance(result, str)
651
661
  assert len(result) > 0
@@ -656,16 +666,15 @@ def test_add_parent_to_baggage():
656
666
  if not _check_otel_installed():
657
667
  pytest.skip("OpenTelemetry SDK not fully installed, skipping test")
658
668
 
659
- from opentelemetry import baggage, context
660
-
661
669
  from braintrust.otel import add_parent_to_baggage
670
+ from opentelemetry import baggage, context
662
671
 
663
672
  # Test adding parent to baggage
664
673
  token = add_parent_to_baggage("project_name:test-project")
665
674
  assert token is not None
666
675
 
667
676
  # Verify it's in baggage
668
- parent_value = baggage.get_baggage('braintrust.parent')
677
+ parent_value = baggage.get_baggage("braintrust.parent")
669
678
  assert parent_value == "project_name:test-project"
670
679
 
671
680
  # Clean up
@@ -677,11 +686,10 @@ def test_add_span_parent_to_baggage():
677
686
  if not _check_otel_installed():
678
687
  pytest.skip("OpenTelemetry SDK not fully installed, skipping test")
679
688
 
689
+ from braintrust.otel import add_span_parent_to_baggage
680
690
  from opentelemetry import baggage, context, trace
681
691
  from opentelemetry.sdk.trace import TracerProvider
682
692
 
683
- from braintrust.otel import add_span_parent_to_baggage
684
-
685
693
  # Setup tracer
686
694
  provider = TracerProvider()
687
695
  trace.set_tracer_provider(provider)
@@ -695,7 +703,7 @@ def test_add_span_parent_to_baggage():
695
703
  assert token is not None
696
704
 
697
705
  # Verify it's in baggage
698
- parent_value = baggage.get_baggage('braintrust.parent')
706
+ parent_value = baggage.get_baggage("braintrust.parent")
699
707
  assert parent_value == "project_name:test"
700
708
 
701
709
  context.detach(token)
braintrust/test_queue.py CHANGED
@@ -3,7 +3,6 @@ import threading
3
3
  import time
4
4
 
5
5
  import pytest
6
-
7
6
  from braintrust.queue import DEFAULT_QUEUE_SIZE, LogQueue
8
7
 
9
8
 
braintrust/test_score.py CHANGED
@@ -65,9 +65,7 @@ class TestScore(unittest.TestCase):
65
65
 
66
66
  def test_from_dict_round_trip(self):
67
67
  """Test that Score can be serialized to dict and deserialized back."""
68
- original = Score(
69
- name="round_trip_scorer", score=0.95, metadata={"info": "test"}
70
- )
68
+ original = Score(name="round_trip_scorer", score=0.95, metadata={"info": "test"})
71
69
 
72
70
  # Serialize to dict
73
71
  as_dict = original.as_dict()