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.
- braintrust/__init__.py +4 -0
- braintrust/_generated_types.py +1200 -611
- braintrust/audit.py +2 -2
- braintrust/cli/eval.py +6 -7
- braintrust/cli/push.py +11 -11
- braintrust/conftest.py +1 -0
- braintrust/context.py +12 -17
- braintrust/contrib/temporal/__init__.py +16 -27
- braintrust/contrib/temporal/test_temporal.py +8 -3
- braintrust/devserver/auth.py +8 -8
- braintrust/devserver/cache.py +3 -4
- braintrust/devserver/cors.py +8 -7
- braintrust/devserver/dataset.py +3 -5
- braintrust/devserver/eval_hooks.py +7 -6
- braintrust/devserver/schemas.py +22 -19
- braintrust/devserver/server.py +19 -12
- braintrust/devserver/test_cached_login.py +4 -4
- braintrust/framework.py +128 -140
- braintrust/framework2.py +88 -87
- braintrust/functions/invoke.py +93 -53
- braintrust/functions/stream.py +3 -2
- braintrust/generated_types.py +17 -1
- braintrust/git_fields.py +11 -11
- braintrust/gitutil.py +2 -3
- braintrust/graph_util.py +10 -10
- braintrust/id_gen.py +2 -2
- braintrust/logger.py +346 -357
- braintrust/merge_row_batch.py +10 -9
- braintrust/oai.py +107 -24
- braintrust/otel/__init__.py +49 -49
- braintrust/otel/context.py +16 -30
- braintrust/otel/test_distributed_tracing.py +14 -11
- braintrust/otel/test_otel_bt_integration.py +32 -31
- braintrust/parameters.py +8 -8
- braintrust/prompt.py +14 -14
- braintrust/prompt_cache/disk_cache.py +5 -4
- braintrust/prompt_cache/lru_cache.py +3 -2
- braintrust/prompt_cache/prompt_cache.py +13 -14
- braintrust/queue.py +4 -4
- braintrust/score.py +4 -4
- braintrust/serializable_data_class.py +4 -4
- braintrust/span_identifier_v1.py +1 -2
- braintrust/span_identifier_v2.py +3 -4
- braintrust/span_identifier_v3.py +23 -20
- braintrust/span_identifier_v4.py +34 -25
- braintrust/test_framework.py +16 -6
- braintrust/test_helpers.py +5 -5
- braintrust/test_id_gen.py +2 -3
- braintrust/test_otel.py +61 -53
- braintrust/test_queue.py +0 -1
- braintrust/test_score.py +1 -3
- braintrust/test_span_components.py +29 -44
- braintrust/util.py +9 -8
- braintrust/version.py +2 -2
- braintrust/wrappers/_anthropic_utils.py +4 -4
- braintrust/wrappers/agno/__init__.py +3 -4
- braintrust/wrappers/agno/agent.py +1 -2
- braintrust/wrappers/agno/function_call.py +1 -2
- braintrust/wrappers/agno/model.py +1 -2
- braintrust/wrappers/agno/team.py +1 -2
- braintrust/wrappers/agno/utils.py +12 -12
- braintrust/wrappers/anthropic.py +7 -8
- braintrust/wrappers/claude_agent_sdk/__init__.py +3 -4
- braintrust/wrappers/claude_agent_sdk/_wrapper.py +29 -27
- braintrust/wrappers/dspy.py +15 -17
- braintrust/wrappers/google_genai/__init__.py +16 -16
- braintrust/wrappers/langchain.py +22 -24
- braintrust/wrappers/litellm.py +4 -3
- braintrust/wrappers/openai.py +15 -15
- braintrust/wrappers/pydantic_ai.py +1204 -0
- braintrust/wrappers/test_agno.py +0 -1
- braintrust/wrappers/test_dspy.py +0 -1
- braintrust/wrappers/test_google_genai.py +2 -3
- braintrust/wrappers/test_litellm.py +0 -1
- braintrust/wrappers/test_oai_attachments.py +322 -0
- braintrust/wrappers/test_pydantic_ai_integration.py +1788 -0
- braintrust/wrappers/{test_pydantic_ai.py → test_pydantic_ai_wrap_openai.py} +1 -2
- {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/METADATA +3 -2
- braintrust-0.4.0.dist-info/RECORD +120 -0
- braintrust-0.3.14.dist-info/RECORD +0 -117
- {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/WHEEL +0 -0
- {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/entry_points.txt +0 -0
- {braintrust-0.3.14.dist-info → braintrust-0.4.0.dist-info}/top_level.txt +0 -0
braintrust/merge_row_batch.py
CHANGED
|
@@ -1,10 +1,11 @@
|
|
|
1
|
-
from
|
|
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 =
|
|
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:
|
|
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:
|
|
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[
|
|
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:
|
|
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=
|
|
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:
|
|
152
|
-
) ->
|
|
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
|
|
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:
|
|
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:
|
|
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":
|
|
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:
|
|
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:
|
|
281
|
+
tool_calls: list[Any] | None = None
|
|
206
282
|
finish_reason = None
|
|
207
|
-
metrics:
|
|
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:
|
|
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:
|
|
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":
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
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
|
-
|
|
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":
|
|
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:
|
|
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:
|
|
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:
|
|
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) ->
|
|
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:
|
|
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) ->
|
|
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
|
braintrust/otel/__init__.py
CHANGED
|
@@ -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:
|
|
142
|
-
api_key:
|
|
143
|
-
parent:
|
|
144
|
-
headers:
|
|
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(
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
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:
|
|
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:
|
|
229
|
-
parent:
|
|
230
|
-
api_url:
|
|
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
|
|
233
|
-
headers:
|
|
234
|
-
SpanProcessor:
|
|
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(
|
|
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(
|
|
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(
|
|
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,
|
|
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(
|
|
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(
|
|
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,
|
|
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(
|
|
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:
|
|
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,
|
|
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,
|
|
578
|
-
span_id_hex = format(span_context.span_id,
|
|
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 ==
|
|
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 ==
|
|
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(
|
|
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(
|
|
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(
|
|
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(
|
|
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(
|