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
@@ -1,10 +1,11 @@
1
- from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
1
+ from collections.abc import Mapping, Sequence
2
+ from typing import Any, Optional
2
3
 
3
4
  from .db_fields import IS_MERGE_FIELD, PARENT_ID_FIELD
4
5
  from .graph_util import UndirectedGraph, topological_sort, undirected_connected_components
5
6
  from .util import merge_dicts
6
7
 
7
- _MergedRowKey = Tuple[Optional[Any], ...]
8
+ _MergedRowKey = tuple[Optional[Any], ...]
8
9
 
9
10
 
10
11
  def _generate_merged_row_key(row: Mapping[str, Any], use_parent_id_for_id: bool = False) -> _MergedRowKey:
@@ -33,7 +34,7 @@ MERGE_ROW_SKIP_FIELDS = [
33
34
  ]
34
35
 
35
36
 
36
- def _pop_merge_row_skip_fields(row: Dict[str, Any]) -> Dict[str, Any]:
37
+ def _pop_merge_row_skip_fields(row: dict[str, Any]) -> dict[str, Any]:
37
38
  popped = {}
38
39
  for field in MERGE_ROW_SKIP_FIELDS:
39
40
  if field in row:
@@ -41,14 +42,14 @@ def _pop_merge_row_skip_fields(row: Dict[str, Any]) -> Dict[str, Any]:
41
42
  return popped
42
43
 
43
44
 
44
- def _restore_merge_row_skip_fields(row: Dict[str, Any], skip_fields: Dict[str, Any]):
45
+ def _restore_merge_row_skip_fields(row: dict[str, Any], skip_fields: dict[str, Any]):
45
46
  for field in MERGE_ROW_SKIP_FIELDS:
46
47
  row.pop(field, None)
47
48
  if field in skip_fields:
48
49
  row[field] = skip_fields[field]
49
50
 
50
51
 
51
- def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
52
+ def merge_row_batch(rows: Sequence[dict[str, Any]]) -> list[list[dict[str, Any]]]:
52
53
  """Given a batch of rows, merges conflicting rows together to end up with a
53
54
  set of rows to insert. Returns a set of de-conflicted rows, as a list of
54
55
  lists, where separate lists contain "independent" rows which can be
@@ -98,7 +99,7 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
98
99
  "Logged row is missing an id. This is an internal braintrust error. Please contact us at info@braintrust.dev for help"
99
100
  )
100
101
 
101
- row_groups: Dict[_MergedRowKey, Dict[str, Any]] = {}
102
+ row_groups: dict[_MergedRowKey, dict[str, Any]] = {}
102
103
  for row in rows:
103
104
  key = _generate_merged_row_key(row)
104
105
  existing_row = row_groups.get(key)
@@ -138,7 +139,7 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
138
139
  # all groups of rows which each row in a group has a PARENT_ID_FIELD
139
140
  # relationship with at least one other row in the group.
140
141
  connected_components = undirected_connected_components(
141
- UndirectedGraph(vertices=set(graph.keys()), edges=set((k, v) for k, vs in graph.items() for v in vs))
142
+ UndirectedGraph(vertices=set(graph.keys()), edges={(k, v) for k, vs in graph.items() for v in vs})
142
143
  )
143
144
 
144
145
  # For each connected row group, run topological sort over that subgraph to
@@ -148,8 +149,8 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
148
149
 
149
150
 
150
151
  def batch_items(
151
- items: List[List[str]], batch_max_num_items: Optional[int] = None, batch_max_num_bytes: Optional[int] = None
152
- ) -> List[List[List[str]]]:
152
+ items: list[list[str]], batch_max_num_items: int | None = None, batch_max_num_bytes: int | None = None
153
+ ) -> list[list[list[str]]]:
153
154
  """Repartition the given list of items into sets of batches which can be
154
155
  published in parallel or in sequence.
155
156
 
braintrust/oai.py CHANGED
@@ -1,8 +1,11 @@
1
1
  import abc
2
+ import base64
3
+ import re
2
4
  import time
3
- from typing import Any, Callable, Dict, List, Optional
5
+ from collections.abc import Callable
6
+ from typing import Any
4
7
 
5
- from .logger import Span, start_span
8
+ from .logger import Attachment, Span, start_span
6
9
  from .span_types import SpanTypeAttribute
7
10
  from .util import merge_dicts
8
11
 
@@ -68,8 +71,77 @@ def log_headers(response: Any, span: Span):
68
71
  )
69
72
 
70
73
 
74
+ def _convert_data_url_to_attachment(data_url: str, filename: str | None = None) -> Attachment | str:
75
+ """Helper function to convert data URL to an Attachment."""
76
+ data_url_match = re.match(r"^data:([^;]+);base64,(.+)$", data_url)
77
+ if not data_url_match:
78
+ return data_url
79
+
80
+ mime_type, base64_data = data_url_match.groups()
81
+
82
+ try:
83
+ binary_data = base64.b64decode(base64_data)
84
+
85
+ if filename is None:
86
+ extension = mime_type.split("/")[1] if "/" in mime_type else "bin"
87
+ prefix = "image" if mime_type.startswith("image/") else "document"
88
+ filename = f"{prefix}.{extension}"
89
+
90
+ attachment = Attachment(data=binary_data, filename=filename, content_type=mime_type)
91
+
92
+ return attachment
93
+ except Exception:
94
+ return data_url
95
+
96
+
97
+ def _process_attachments_in_input(input_data: Any) -> Any:
98
+ """Process input to convert data URL images and base64 documents to Attachment objects."""
99
+ if isinstance(input_data, list):
100
+ return [_process_attachments_in_input(item) for item in input_data]
101
+
102
+ if isinstance(input_data, dict):
103
+ # Check for OpenAI's image_url format with data URLs
104
+ if (
105
+ input_data.get("type") == "image_url"
106
+ and isinstance(input_data.get("image_url"), dict)
107
+ and isinstance(input_data["image_url"].get("url"), str)
108
+ ):
109
+ processed_url = _convert_data_url_to_attachment(input_data["image_url"]["url"])
110
+ return {
111
+ **input_data,
112
+ "image_url": {
113
+ **input_data["image_url"],
114
+ "url": processed_url,
115
+ },
116
+ }
117
+
118
+ # Check for OpenAI's file format with data URL (e.g., PDFs)
119
+ if (
120
+ input_data.get("type") == "file"
121
+ and isinstance(input_data.get("file"), dict)
122
+ and isinstance(input_data["file"].get("file_data"), str)
123
+ ):
124
+ file_filename = input_data["file"].get("filename")
125
+ processed_file_data = _convert_data_url_to_attachment(
126
+ input_data["file"]["file_data"],
127
+ filename=file_filename if isinstance(file_filename, str) else None,
128
+ )
129
+ return {
130
+ **input_data,
131
+ "file": {
132
+ **input_data["file"],
133
+ "file_data": processed_file_data,
134
+ },
135
+ }
136
+
137
+ # Recursively process nested objects
138
+ return {key: _process_attachments_in_input(value) for key, value in input_data.items()}
139
+
140
+ return input_data
141
+
142
+
71
143
  class ChatCompletionWrapper:
72
- def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
144
+ def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
73
145
  self.create_fn = create_fn
74
146
  self.acreate_fn = acreate_fn
75
147
 
@@ -183,28 +255,32 @@ class ChatCompletionWrapper:
183
255
  span.end()
184
256
 
185
257
  @classmethod
186
- def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
258
+ def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
187
259
  # First, destructively remove span_info
188
260
  ret = params.pop("span_info", {})
189
261
 
190
262
  # Then, copy the rest of the params
191
263
  params = prettify_params(params)
192
264
  messages = params.pop("messages", None)
265
+
266
+ # Process attachments in input (convert data URLs to Attachment objects)
267
+ processed_input = _process_attachments_in_input(messages)
268
+
193
269
  return merge_dicts(
194
270
  ret,
195
271
  {
196
- "input": messages,
272
+ "input": processed_input,
197
273
  "metadata": {**params, "provider": "openai"},
198
274
  },
199
275
  )
200
276
 
201
277
  @classmethod
202
- def _postprocess_streaming_results(cls, all_results: List[Dict[str, Any]]) -> Dict[str, Any]:
278
+ def _postprocess_streaming_results(cls, all_results: list[dict[str, Any]]) -> dict[str, Any]:
203
279
  role = None
204
280
  content = None
205
- tool_calls: Optional[List[Any]] = None
281
+ tool_calls: list[Any] | None = None
206
282
  finish_reason = None
207
- metrics: Dict[str, float] = {}
283
+ metrics: dict[str, float] = {}
208
284
  for result in all_results:
209
285
  usage = result.get("usage")
210
286
  if usage:
@@ -263,7 +339,7 @@ class ChatCompletionWrapper:
263
339
 
264
340
 
265
341
  class ResponseWrapper:
266
- def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]], name: str = "openai.responses.create"):
342
+ def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None, name: str = "openai.responses.create"):
267
343
  self.create_fn = create_fn
268
344
  self.acreate_fn = acreate_fn
269
345
  self.name = name
@@ -372,23 +448,27 @@ class ResponseWrapper:
372
448
  span.end()
373
449
 
374
450
  @classmethod
375
- def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
451
+ def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
376
452
  # First, destructively remove span_info
377
453
  ret = params.pop("span_info", {})
378
454
 
379
455
  # Then, copy the rest of the params
380
456
  params = prettify_params(params)
381
457
  input_data = params.pop("input", None)
458
+
459
+ # Process attachments in input (convert data URLs to Attachment objects)
460
+ processed_input = _process_attachments_in_input(input_data)
461
+
382
462
  return merge_dicts(
383
463
  ret,
384
464
  {
385
- "input": input_data,
465
+ "input": processed_input,
386
466
  "metadata": {**params, "provider": "openai"},
387
467
  },
388
468
  )
389
469
 
390
470
  @classmethod
391
- def _parse_event_from_result(cls, result: Dict[str, Any]) -> Dict[str, Any]:
471
+ def _parse_event_from_result(cls, result: dict[str, Any]) -> dict[str, Any]:
392
472
  """Parse event from response result"""
393
473
  data = {"metrics": {}}
394
474
 
@@ -408,7 +488,7 @@ class ResponseWrapper:
408
488
  return data
409
489
 
410
490
  @classmethod
411
- def _postprocess_streaming_results(cls, all_results: List[Any]) -> Dict[str, Any]:
491
+ def _postprocess_streaming_results(cls, all_results: list[Any]) -> dict[str, Any]:
412
492
  """Process streaming results - minimal version focused on metrics extraction."""
413
493
  metrics = {}
414
494
  output = []
@@ -491,13 +571,13 @@ class ResponseWrapper:
491
571
 
492
572
 
493
573
  class BaseWrapper(abc.ABC):
494
- def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]], name: str):
574
+ def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None, name: str):
495
575
  self._create_fn = create_fn
496
576
  self._acreate_fn = acreate_fn
497
577
  self._name = name
498
578
 
499
579
  @abc.abstractmethod
500
- def process_output(self, response: Dict[str, Any], span: Span):
580
+ def process_output(self, response: dict[str, Any], span: Span):
501
581
  """Process the API response and log relevant information to the span."""
502
582
  pass
503
583
 
@@ -535,27 +615,30 @@ class BaseWrapper(abc.ABC):
535
615
  return raw_response
536
616
 
537
617
  @classmethod
538
- def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
618
+ def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
539
619
  # First, destructively remove span_info
540
620
  ret = params.pop("span_info", {})
541
621
 
542
622
  params = prettify_params(params)
543
- input = params.pop("input", None)
623
+ input_data = params.pop("input", None)
624
+
625
+ # Process attachments in input (convert data URLs to Attachment objects)
626
+ processed_input = _process_attachments_in_input(input_data)
544
627
 
545
628
  return merge_dicts(
546
629
  ret,
547
630
  {
548
- "input": input,
631
+ "input": processed_input,
549
632
  "metadata": {**params, "provider": "openai"},
550
633
  },
551
634
  )
552
635
 
553
636
 
554
637
  class EmbeddingWrapper(BaseWrapper):
555
- def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
638
+ def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
556
639
  super().__init__(create_fn, acreate_fn, "Embedding")
557
640
 
558
- def process_output(self, response: Dict[str, Any], span: Span):
641
+ def process_output(self, response: dict[str, Any], span: Span):
559
642
  usage = response.get("usage")
560
643
  metrics = _parse_metrics_from_usage(usage)
561
644
  span.log(
@@ -567,7 +650,7 @@ class EmbeddingWrapper(BaseWrapper):
567
650
 
568
651
 
569
652
  class ModerationWrapper(BaseWrapper):
570
- def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
653
+ def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
571
654
  super().__init__(create_fn, acreate_fn, "Moderation")
572
655
 
573
656
  def process_output(self, response: Any, span: Span):
@@ -814,7 +897,7 @@ TOKEN_PREFIX_MAP = {
814
897
  }
815
898
 
816
899
 
817
- def _parse_metrics_from_usage(usage: Any) -> Dict[str, Any]:
900
+ def _parse_metrics_from_usage(usage: Any) -> dict[str, Any]:
818
901
  # For simplicity, this function handles all the different APIs
819
902
  metrics = {}
820
903
 
@@ -848,7 +931,7 @@ def _is_numeric(v):
848
931
  return isinstance(v, (int, float, complex)) and not isinstance(v, bool)
849
932
 
850
933
 
851
- def prettify_params(params: Dict[str, Any]) -> Dict[str, Any]:
934
+ def prettify_params(params: dict[str, Any]) -> dict[str, Any]:
852
935
  # Filter out NOT_GIVEN parameters
853
936
  # https://linear.app/braintrustdata/issue/BRA-2467
854
937
  ret = {k: v for k, v in params.items() if not _is_not_given(v)}
@@ -858,7 +941,7 @@ def prettify_params(params: Dict[str, Any]) -> Dict[str, Any]:
858
941
  return ret
859
942
 
860
943
 
861
- def _try_to_dict(obj: Any) -> Dict[str, Any]:
944
+ def _try_to_dict(obj: Any) -> dict[str, Any]:
862
945
  if isinstance(obj, dict):
863
946
  return obj
864
947
  # convert a pydantic object to a dict
@@ -1,7 +1,6 @@
1
1
  import logging
2
2
  import os
3
3
  import warnings
4
- from typing import Dict, Optional
5
4
  from urllib.parse import urljoin
6
5
 
7
6
  INSTALL_ERR_MSG = (
@@ -138,10 +137,10 @@ class OtelExporter(OTLPSpanExporter):
138
137
 
139
138
  def __init__(
140
139
  self,
141
- url: Optional[str] = None,
142
- api_key: Optional[str] = None,
143
- parent: Optional[str] = None,
144
- headers: Optional[Dict[str, str]] = None,
140
+ url: str | None = None,
141
+ api_key: str | None = None,
142
+ parent: str | None = None,
143
+ headers: dict[str, str] | None = None,
145
144
  **kwargs,
146
145
  ):
147
146
  """
@@ -189,13 +188,14 @@ class OtelExporter(OTLPSpanExporter):
189
188
  super().__init__(endpoint=endpoint, headers=exporter_headers, **kwargs)
190
189
 
191
190
 
192
- def add_braintrust_span_processor(tracer_provider,
193
- api_key: Optional[str] = None,
194
- parent: Optional[str] = None,
195
- api_url: Optional[str] = None,
191
+ def add_braintrust_span_processor(
192
+ tracer_provider,
193
+ api_key: str | None = None,
194
+ parent: str | None = None,
195
+ api_url: str | None = None,
196
196
  filter_ai_spans: bool = False,
197
197
  custom_filter=None,
198
- headers: Optional[Dict[str, str]] = None,
198
+ headers: dict[str, str] | None = None,
199
199
  ):
200
200
  processor = BraintrustSpanProcessor(
201
201
  api_key=api_key,
@@ -225,13 +225,13 @@ class BraintrustSpanProcessor:
225
225
 
226
226
  def __init__(
227
227
  self,
228
- api_key: Optional[str] = None,
229
- parent: Optional[str] = None,
230
- api_url: Optional[str] = None,
228
+ api_key: str | None = None,
229
+ parent: str | None = None,
230
+ api_url: str | None = None,
231
231
  filter_ai_spans: bool = False,
232
- custom_filter = None,
233
- headers: Optional[Dict[str, str]] = None,
234
- SpanProcessor: Optional[type] = None,
232
+ custom_filter=None,
233
+ headers: dict[str, str] | None = None,
234
+ SpanProcessor: type | None = None,
235
235
  ):
236
236
  """
237
237
  Initialize the BraintrustSpanProcessor.
@@ -279,16 +279,17 @@ class BraintrustSpanProcessor:
279
279
 
280
280
  # Priority 1: Check if braintrust.parent is in current OTEL context
281
281
  from opentelemetry import baggage, context
282
+
282
283
  current_context = context.get_current()
283
- parent_value = context.get_value('braintrust.parent', current_context)
284
+ parent_value = context.get_value("braintrust.parent", current_context)
284
285
 
285
286
  # Priority 2: Check OTEL baggage (propagates automatically across contexts)
286
287
  if not parent_value:
287
- parent_value = baggage.get_baggage('braintrust.parent', context=current_context)
288
+ parent_value = baggage.get_baggage("braintrust.parent", context=current_context)
288
289
 
289
290
  # Priority 3: Check if parent_context has braintrust.parent (backup)
290
291
  if not parent_value and parent_context:
291
- parent_value = context.get_value('braintrust.parent', parent_context)
292
+ parent_value = context.get_value("braintrust.parent", parent_context)
292
293
 
293
294
  # Priority 4: Check if parent OTEL span has braintrust.parent attribute
294
295
  if not parent_value and parent_context:
@@ -304,7 +305,6 @@ class BraintrustSpanProcessor:
304
305
 
305
306
  self._processor.on_start(span, parent_context)
306
307
 
307
-
308
308
  def _get_parent_otel_braintrust_parent(self, parent_context):
309
309
  """Get braintrust.parent attribute from parent OTEL span if it exists."""
310
310
  try:
@@ -313,7 +313,7 @@ class BraintrustSpanProcessor:
313
313
  # Get the current span from the parent context
314
314
  current_span = trace.get_current_span(parent_context)
315
315
 
316
- if current_span and hasattr(current_span, 'attributes') and current_span.attributes:
316
+ if current_span and hasattr(current_span, "attributes") and current_span.attributes:
317
317
  # Check if parent span has braintrust.parent attribute
318
318
  attributes = dict(current_span.attributes)
319
319
  return attributes.get("braintrust.parent")
@@ -346,11 +346,7 @@ class BraintrustSpanProcessor:
346
346
  return self._processor
347
347
 
348
348
 
349
- def _get_braintrust_parent(
350
- object_type,
351
- object_id: Optional[str] = None,
352
- compute_args: Optional[Dict] = None
353
- ) -> Optional[str]:
349
+ def _get_braintrust_parent(object_type, object_id: str | None = None, compute_args: dict | None = None) -> str | None:
354
350
  """
355
351
  Construct a braintrust.parent identifier string from span components.
356
352
 
@@ -405,11 +401,10 @@ def context_from_span_export(export_str: str):
405
401
  if not OTEL_AVAILABLE:
406
402
  raise ImportError(INSTALL_ERR_MSG)
407
403
 
404
+ from braintrust.span_identifier_v4 import SpanComponentsV4
408
405
  from opentelemetry import baggage, trace
409
406
  from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
410
407
 
411
- from braintrust.span_identifier_v4 import SpanComponentsV4
412
-
413
408
  # Parse the export string (handles V3/V4 automatically)
414
409
  components = SpanComponentsV4.from_str(export_str)
415
410
 
@@ -417,7 +412,7 @@ def context_from_span_export(export_str: str):
417
412
  braintrust_parent = _get_braintrust_parent(
418
413
  object_type=components.object_type,
419
414
  object_id=components.object_id,
420
- compute_args=components.compute_object_metadata_args
415
+ compute_args=components.compute_object_metadata_args,
421
416
  )
422
417
 
423
418
  # Convert hex strings to OTEL integers
@@ -429,7 +424,7 @@ def context_from_span_export(export_str: str):
429
424
  trace_id=trace_id_int,
430
425
  span_id=span_id_int,
431
426
  is_remote=True, # Critical: mark as remote for distributed tracing
432
- trace_flags=TraceFlags(TraceFlags.SAMPLED)
427
+ trace_flags=TraceFlags(TraceFlags.SAMPLED),
433
428
  )
434
429
 
435
430
  # Create NonRecordingSpan and set in context
@@ -438,7 +433,7 @@ def context_from_span_export(export_str: str):
438
433
 
439
434
  # Set braintrust.parent in OTEL baggage so it propagates automatically
440
435
  if braintrust_parent:
441
- ctx = baggage.set_baggage('braintrust.parent', braintrust_parent, context=ctx)
436
+ ctx = baggage.set_baggage("braintrust.parent", braintrust_parent, context=ctx)
442
437
 
443
438
  return ctx
444
439
 
@@ -475,7 +470,7 @@ def add_parent_to_baggage(parent: str, ctx=None):
475
470
  from opentelemetry import baggage, context
476
471
 
477
472
  # Set in baggage so it propagates via inject()
478
- new_ctx = baggage.set_baggage('braintrust.parent', parent, context=ctx)
473
+ new_ctx = baggage.set_baggage("braintrust.parent", parent, context=ctx)
479
474
  token = context.attach(new_ctx)
480
475
  return token
481
476
 
@@ -511,11 +506,11 @@ def add_span_parent_to_baggage(span, ctx=None):
511
506
  raise ImportError(INSTALL_ERR_MSG)
512
507
 
513
508
  # Get braintrust.parent from span attributes
514
- if not span or not hasattr(span, 'attributes') or not span.attributes:
509
+ if not span or not hasattr(span, "attributes") or not span.attributes:
515
510
  logging.warning("add_span_parent_to_baggage: span has no attributes")
516
511
  return None
517
512
 
518
- parent_value = span.attributes.get('braintrust.parent')
513
+ parent_value = span.attributes.get("braintrust.parent")
519
514
  if not parent_value:
520
515
  logging.warning(
521
516
  "add_span_parent_to_baggage: braintrust.parent attribute not found. "
@@ -527,7 +522,7 @@ def add_span_parent_to_baggage(span, ctx=None):
527
522
  return add_parent_to_baggage(parent_value, ctx=ctx)
528
523
 
529
524
 
530
- def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
525
+ def parent_from_headers(headers: dict[str, str]) -> str | None:
531
526
  """
532
527
  Extract a Braintrust-compatible parent string from W3C Trace Context headers.
533
528
 
@@ -554,17 +549,16 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
554
549
  if not OTEL_AVAILABLE:
555
550
  raise ImportError(INSTALL_ERR_MSG)
556
551
 
552
+ from braintrust.span_identifier_v4 import SpanComponentsV4
557
553
  from opentelemetry import baggage, trace
558
554
  from opentelemetry.propagate import extract
559
555
 
560
- from braintrust.span_identifier_v4 import SpanComponentsV4
561
-
562
556
  # Extract context from headers using W3C Trace Context propagator
563
557
  ctx = extract(headers)
564
558
 
565
559
  # Get span from context
566
560
  span = trace.get_current_span(ctx)
567
- if not span or not hasattr(span, 'get_span_context'):
561
+ if not span or not hasattr(span, "get_span_context"):
568
562
  logging.error("parent_from_headers: No valid span found in headers")
569
563
  return None
570
564
 
@@ -574,19 +568,19 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
574
568
  return None
575
569
 
576
570
  # Convert OTEL IDs to hex strings
577
- trace_id_hex = format(span_context.trace_id, '032x')
578
- span_id_hex = format(span_context.span_id, '016x')
571
+ trace_id_hex = format(span_context.trace_id, "032x")
572
+ span_id_hex = format(span_context.span_id, "016x")
579
573
 
580
574
  # Validate trace_id and span_id are not all zeros
581
- if trace_id_hex == '00000000000000000000000000000000':
575
+ if trace_id_hex == "00000000000000000000000000000000":
582
576
  logging.error("parent_from_headers: Invalid trace_id (all zeros)")
583
577
  return None
584
- if span_id_hex == '0000000000000000':
578
+ if span_id_hex == "0000000000000000":
585
579
  logging.error("parent_from_headers: Invalid span_id (all zeros)")
586
580
  return None
587
581
 
588
582
  # Get braintrust.parent from baggage if present
589
- braintrust_parent = baggage.get_baggage('braintrust.parent', context=ctx)
583
+ braintrust_parent = baggage.get_baggage("braintrust.parent", context=ctx)
590
584
 
591
585
  # Parse braintrust.parent to extract object_type and object_id
592
586
  object_type = None
@@ -607,22 +601,28 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
607
601
  # Parse braintrust.parent format: "project_id:abc", "project_name:xyz", or "experiment_id:123"
608
602
  if braintrust_parent.startswith("project_id:"):
609
603
  object_type = SpanObjectTypeV3.PROJECT_LOGS
610
- object_id = braintrust_parent[len("project_id:"):]
604
+ object_id = braintrust_parent[len("project_id:") :]
611
605
  if not object_id:
612
- logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty project_id): {braintrust_parent}")
606
+ logging.error(
607
+ f"parent_from_headers: Invalid braintrust.parent format (empty project_id): {braintrust_parent}"
608
+ )
613
609
  return None
614
610
  elif braintrust_parent.startswith("project_name:"):
615
611
  object_type = SpanObjectTypeV3.PROJECT_LOGS
616
- project_name = braintrust_parent[len("project_name:"):]
612
+ project_name = braintrust_parent[len("project_name:") :]
617
613
  if not project_name:
618
- logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty project_name): {braintrust_parent}")
614
+ logging.error(
615
+ f"parent_from_headers: Invalid braintrust.parent format (empty project_name): {braintrust_parent}"
616
+ )
619
617
  return None
620
618
  compute_args = {"project_name": project_name}
621
619
  elif braintrust_parent.startswith("experiment_id:"):
622
620
  object_type = SpanObjectTypeV3.EXPERIMENT
623
- object_id = braintrust_parent[len("experiment_id:"):]
621
+ object_id = braintrust_parent[len("experiment_id:") :]
624
622
  if not object_id:
625
- logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty experiment_id): {braintrust_parent}")
623
+ logging.error(
624
+ f"parent_from_headers: Invalid braintrust.parent format (empty experiment_id): {braintrust_parent}"
625
+ )
626
626
  return None
627
627
  else:
628
628
  logging.error(