braintrust 0.3.8__py3-none-any.whl → 0.3.10__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/_generated_types.py +52 -9
- braintrust/cli/install/api.py +0 -24
- braintrust/cli/push.py +2 -1
- braintrust/conftest.py +9 -0
- braintrust/contrib/__init__.py +5 -0
- braintrust/contrib/temporal/__init__.py +438 -0
- braintrust/contrib/temporal/test_temporal.py +502 -0
- braintrust/framework2.py +19 -1
- braintrust/generated_types.py +5 -1
- braintrust/logger.py +20 -9
- braintrust/test_framework2.py +233 -0
- braintrust/test_logger.py +43 -1
- braintrust/version.py +2 -2
- {braintrust-0.3.8.dist-info → braintrust-0.3.10.dist-info}/METADATA +4 -1
- {braintrust-0.3.8.dist-info → braintrust-0.3.10.dist-info}/RECORD +18 -14
- {braintrust-0.3.8.dist-info → braintrust-0.3.10.dist-info}/WHEEL +0 -0
- {braintrust-0.3.8.dist-info → braintrust-0.3.10.dist-info}/entry_points.txt +0 -0
- {braintrust-0.3.8.dist-info → braintrust-0.3.10.dist-info}/top_level.txt +0 -0
braintrust/_generated_types.py
CHANGED
|
@@ -212,6 +212,17 @@ CallEvent = Union[
|
|
|
212
212
|
]
|
|
213
213
|
|
|
214
214
|
|
|
215
|
+
class ChatCompletionContentPartFileFile(TypedDict):
|
|
216
|
+
file_data: NotRequired[Optional[str]]
|
|
217
|
+
filename: NotRequired[Optional[str]]
|
|
218
|
+
file_id: NotRequired[Optional[str]]
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
class ChatCompletionContentPartFileWithTitle(TypedDict):
|
|
222
|
+
file: ChatCompletionContentPartFileFile
|
|
223
|
+
type: Literal['file']
|
|
224
|
+
|
|
225
|
+
|
|
215
226
|
class ChatCompletionContentPartImageWithTitleImageUrl(TypedDict):
|
|
216
227
|
url: str
|
|
217
228
|
detail: NotRequired[Optional[Union[Literal['auto'], Literal['low'], Literal['high']]]]
|
|
@@ -342,7 +353,7 @@ class ChatCompletionTool(TypedDict):
|
|
|
342
353
|
|
|
343
354
|
|
|
344
355
|
class CodeBundleRuntimeContext(TypedDict):
|
|
345
|
-
runtime: Literal['node', 'python']
|
|
356
|
+
runtime: Literal['node', 'python', 'browser']
|
|
346
357
|
version: str
|
|
347
358
|
|
|
348
359
|
|
|
@@ -570,7 +581,7 @@ class Data(CodeBundle):
|
|
|
570
581
|
|
|
571
582
|
|
|
572
583
|
class FunctionDataFunctionData1DataRuntimeContext(TypedDict):
|
|
573
|
-
runtime: Literal['node', 'python']
|
|
584
|
+
runtime: Literal['node', 'python', 'browser']
|
|
574
585
|
version: str
|
|
575
586
|
|
|
576
587
|
|
|
@@ -649,7 +660,7 @@ class FunctionIdFunctionId3(TypedDict):
|
|
|
649
660
|
|
|
650
661
|
|
|
651
662
|
class FunctionIdFunctionId4InlineContext(TypedDict):
|
|
652
|
-
runtime: Literal['node', 'python']
|
|
663
|
+
runtime: Literal['node', 'python', 'browser']
|
|
653
664
|
version: str
|
|
654
665
|
|
|
655
666
|
|
|
@@ -668,16 +679,16 @@ class FunctionIdFunctionId4(TypedDict):
|
|
|
668
679
|
FunctionIdRef = Mapping[str, Any]
|
|
669
680
|
|
|
670
681
|
|
|
671
|
-
FunctionObjectType = Literal['prompt', 'tool', 'scorer', 'task', 'agent']
|
|
682
|
+
FunctionObjectType = Literal['prompt', 'tool', 'scorer', 'task', 'agent', 'custom_view']
|
|
672
683
|
|
|
673
684
|
|
|
674
685
|
FunctionOutputType = Literal['completion', 'score', 'any']
|
|
675
686
|
|
|
676
687
|
|
|
677
|
-
FunctionTypeEnum = Literal['llm', 'scorer', 'task', 'tool']
|
|
688
|
+
FunctionTypeEnum = Literal['llm', 'scorer', 'task', 'tool', 'custom_view']
|
|
678
689
|
|
|
679
690
|
|
|
680
|
-
FunctionTypeEnumNullish = Literal['llm', 'scorer', 'task', 'tool']
|
|
691
|
+
FunctionTypeEnumNullish = Literal['llm', 'scorer', 'task', 'tool', 'custom_view']
|
|
681
692
|
|
|
682
693
|
|
|
683
694
|
class GitMetadataSettings(TypedDict):
|
|
@@ -1854,7 +1865,7 @@ class AnyModelParams(TypedDict):
|
|
|
1854
1865
|
function_call: NotRequired[Optional[Union[Literal['auto'], Literal['none'], AnyModelParamsFunctionCall]]]
|
|
1855
1866
|
n: NotRequired[Optional[float]]
|
|
1856
1867
|
stop: NotRequired[Optional[Sequence[str]]]
|
|
1857
|
-
reasoning_effort: NotRequired[Optional[Literal['minimal', 'low', 'medium', 'high']]]
|
|
1868
|
+
reasoning_effort: NotRequired[Optional[Literal['none', 'minimal', 'low', 'medium', 'high']]]
|
|
1858
1869
|
verbosity: NotRequired[Optional[Literal['low', 'medium', 'high']]]
|
|
1859
1870
|
top_k: NotRequired[Optional[float]]
|
|
1860
1871
|
stop_sequences: NotRequired[Optional[Sequence[str]]]
|
|
@@ -1894,7 +1905,11 @@ class AttachmentStatus(TypedDict):
|
|
|
1894
1905
|
"""
|
|
1895
1906
|
|
|
1896
1907
|
|
|
1897
|
-
ChatCompletionContentPart = Union[
|
|
1908
|
+
ChatCompletionContentPart = Union[
|
|
1909
|
+
ChatCompletionContentPartTextWithTitle,
|
|
1910
|
+
ChatCompletionContentPartImageWithTitle,
|
|
1911
|
+
ChatCompletionContentPartFileWithTitle,
|
|
1912
|
+
]
|
|
1898
1913
|
|
|
1899
1914
|
|
|
1900
1915
|
class ChatCompletionMessageParamChatCompletionMessageParam1(TypedDict):
|
|
@@ -1993,6 +2008,14 @@ class DatasetEvent(TypedDict):
|
|
|
1993
2008
|
Whether this span is a root span
|
|
1994
2009
|
"""
|
|
1995
2010
|
origin: NotRequired[Optional[ObjectReferenceNullish]]
|
|
2011
|
+
comments: NotRequired[Optional[Sequence[Any]]]
|
|
2012
|
+
"""
|
|
2013
|
+
Optional list of comments attached to this event
|
|
2014
|
+
"""
|
|
2015
|
+
audit_data: NotRequired[Optional[Sequence[Any]]]
|
|
2016
|
+
"""
|
|
2017
|
+
Optional list of audit entries attached to this event
|
|
2018
|
+
"""
|
|
1996
2019
|
|
|
1997
2020
|
|
|
1998
2021
|
class Experiment(TypedDict):
|
|
@@ -2075,7 +2098,7 @@ class ModelParamsModelParams(TypedDict):
|
|
|
2075
2098
|
function_call: NotRequired[Optional[Union[Literal['auto'], Literal['none'], ModelParamsModelParamsFunctionCall]]]
|
|
2076
2099
|
n: NotRequired[Optional[float]]
|
|
2077
2100
|
stop: NotRequired[Optional[Sequence[str]]]
|
|
2078
|
-
reasoning_effort: NotRequired[Optional[Literal['minimal', 'low', 'medium', 'high']]]
|
|
2101
|
+
reasoning_effort: NotRequired[Optional[Literal['none', 'minimal', 'low', 'medium', 'high']]]
|
|
2079
2102
|
verbosity: NotRequired[Optional[Literal['low', 'medium', 'high']]]
|
|
2080
2103
|
|
|
2081
2104
|
|
|
@@ -2327,6 +2350,14 @@ class ExperimentEvent(TypedDict):
|
|
|
2327
2350
|
Whether this span is a root span
|
|
2328
2351
|
"""
|
|
2329
2352
|
origin: NotRequired[Optional[ObjectReferenceNullish]]
|
|
2353
|
+
comments: NotRequired[Optional[Sequence[Any]]]
|
|
2354
|
+
"""
|
|
2355
|
+
Optional list of comments attached to this event
|
|
2356
|
+
"""
|
|
2357
|
+
audit_data: NotRequired[Optional[Sequence[Any]]]
|
|
2358
|
+
"""
|
|
2359
|
+
Optional list of audit entries attached to this event
|
|
2360
|
+
"""
|
|
2330
2361
|
|
|
2331
2362
|
|
|
2332
2363
|
class GraphNodeGraphNode7(TypedDict):
|
|
@@ -2437,6 +2468,18 @@ class ProjectLogsEvent(TypedDict):
|
|
|
2437
2468
|
"""
|
|
2438
2469
|
span_attributes: NotRequired[Optional[SpanAttributes]]
|
|
2439
2470
|
origin: NotRequired[Optional[ObjectReferenceNullish]]
|
|
2471
|
+
comments: NotRequired[Optional[Sequence[Any]]]
|
|
2472
|
+
"""
|
|
2473
|
+
Optional list of comments attached to this event
|
|
2474
|
+
"""
|
|
2475
|
+
audit_data: NotRequired[Optional[Sequence[Any]]]
|
|
2476
|
+
"""
|
|
2477
|
+
Optional list of audit entries attached to this event
|
|
2478
|
+
"""
|
|
2479
|
+
_async_scoring_state: NotRequired[Optional[Any]]
|
|
2480
|
+
"""
|
|
2481
|
+
The async scoring state for this event
|
|
2482
|
+
"""
|
|
2440
2483
|
|
|
2441
2484
|
|
|
2442
2485
|
class ProjectScore(TypedDict):
|
braintrust/cli/install/api.py
CHANGED
|
@@ -4,7 +4,6 @@ import textwrap
|
|
|
4
4
|
import time
|
|
5
5
|
|
|
6
6
|
from botocore.exceptions import ClientError
|
|
7
|
-
|
|
8
7
|
from braintrust.logger import app_conn, login
|
|
9
8
|
|
|
10
9
|
# pylint: disable=no-name-in-module
|
|
@@ -30,8 +29,6 @@ PARAMS = {
|
|
|
30
29
|
"PrivateSubnet2CIDR": "private_subnet_2_cidr",
|
|
31
30
|
"PrivateSubnet3CIDR": "private_subnet_3_cidr",
|
|
32
31
|
"ManagedPostgres": "managed_postgres",
|
|
33
|
-
"ManagedClickhouse": "managed_clickhouse",
|
|
34
|
-
"ClickhouseInstanceType": "clickhouse_instance_type",
|
|
35
32
|
"PostgresVersion": "postgres_version",
|
|
36
33
|
"OutboundRateLimitWindowMinutes": "outbound_rate_limit_window_minutes",
|
|
37
34
|
"OutboundRateLimitMaxRequests": "outbound_rate_limit_max_requests",
|
|
@@ -179,19 +176,6 @@ def build_parser(subparsers, parents):
|
|
|
179
176
|
default=None,
|
|
180
177
|
)
|
|
181
178
|
|
|
182
|
-
# Clickhouse
|
|
183
|
-
parser.add_argument(
|
|
184
|
-
"--managed-clickhouse",
|
|
185
|
-
help="Spin up a Clickhouse Instance for faster analytics",
|
|
186
|
-
default=None,
|
|
187
|
-
choices=[None, "true", "false"],
|
|
188
|
-
)
|
|
189
|
-
parser.add_argument(
|
|
190
|
-
"--clickhouse-instance-type",
|
|
191
|
-
help="The instance type for the Clickhouse instance",
|
|
192
|
-
default=None,
|
|
193
|
-
)
|
|
194
|
-
|
|
195
179
|
# ElastiCacheClusterId
|
|
196
180
|
parser.add_argument("--elasticache-cluster-host", help="The ElastiCacheCluster host to use", default=None)
|
|
197
181
|
parser.add_argument(
|
|
@@ -236,11 +220,6 @@ def build_parser(subparsers, parents):
|
|
|
236
220
|
help="[Advanced] The postgres URL to use (if you are connecting to another VPC)",
|
|
237
221
|
default=None,
|
|
238
222
|
)
|
|
239
|
-
parser.add_argument("--clickhouse-pg-url", help="[Advanced] The clickhouse PG URL to use", default=None)
|
|
240
|
-
parser.add_argument("--clickhouse-connect-url", help="[Advanced] The clickhouse connect URL to use", default=None)
|
|
241
|
-
parser.add_argument(
|
|
242
|
-
"--clickhouse-catchup-etl-arn", help="[Advanced] The clickhouse catchup ETL ARN to use", default=None
|
|
243
|
-
)
|
|
244
223
|
|
|
245
224
|
# To configure your org
|
|
246
225
|
parser.add_argument(
|
|
@@ -321,9 +300,6 @@ def main(args):
|
|
|
321
300
|
PARAMS["ElastiCacheClusterHost"] = "elasticache_cluster_host"
|
|
322
301
|
PARAMS["ElastiCacheClusterPort"] = "elasticache_cluster_port"
|
|
323
302
|
PARAMS["PostgresUrl"] = "postgres_url"
|
|
324
|
-
PARAMS["ClickhouseCatchupEtlArn"] = "clickhouse_catchup_etl_arn"
|
|
325
|
-
PARAMS["ClickhouseConnectUrl"] = "clickhouse_connect_url"
|
|
326
|
-
PARAMS["ClickhousePGUrl"] = "clickhouse_pg_url"
|
|
327
303
|
PARAMS["EnableBrainstore"] = "enable_brainstore"
|
|
328
304
|
PARAMS["BrainstoreInstanceKeyPairName"] = "brainstore_instance_key_pair_name"
|
|
329
305
|
PARAMS["BrainstoreLicenseKey"] = "brainstore_license_key"
|
braintrust/cli/push.py
CHANGED
|
@@ -16,7 +16,6 @@ import zipfile
|
|
|
16
16
|
from typing import Any, Dict, List, Optional
|
|
17
17
|
|
|
18
18
|
import requests
|
|
19
|
-
|
|
20
19
|
from braintrust.framework import _set_lazy_load
|
|
21
20
|
|
|
22
21
|
from .. import api_conn, login, org_id, proxy_conn
|
|
@@ -250,6 +249,8 @@ def _collect_function_function_defs(
|
|
|
250
249
|
},
|
|
251
250
|
"if_exists": f.if_exists if f.if_exists else if_exists,
|
|
252
251
|
}
|
|
252
|
+
if f.metadata is not None:
|
|
253
|
+
j["metadata"] = f.metadata
|
|
253
254
|
if f.parameters is None:
|
|
254
255
|
raise ValueError(f"Function {f.name} has no supplied parameters")
|
|
255
256
|
j["function_schema"] = {
|
braintrust/conftest.py
CHANGED
|
@@ -36,3 +36,12 @@ def override_app_url_for_tests():
|
|
|
36
36
|
@pytest.fixture(autouse=True)
|
|
37
37
|
def setup_braintrust():
|
|
38
38
|
os.environ.setdefault("GOOGLE_API_KEY", os.getenv("GEMINI_API_KEY", "your_google_api_key_here"))
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
@pytest.fixture(autouse=True)
|
|
42
|
+
def reset_braintrust_state():
|
|
43
|
+
"""Reset all Braintrust global state after each test."""
|
|
44
|
+
yield
|
|
45
|
+
from braintrust import logger
|
|
46
|
+
|
|
47
|
+
logger._state = logger.BraintrustState()
|
|
@@ -0,0 +1,438 @@
|
|
|
1
|
+
"""Braintrust integration for Temporal workflows and activities.
|
|
2
|
+
|
|
3
|
+
This module provides Temporal integration that automatically traces workflow executions
|
|
4
|
+
and activities in Braintrust. To use this integration, install braintrust with the
|
|
5
|
+
temporal extra:
|
|
6
|
+
|
|
7
|
+
pip install braintrust[temporal]
|
|
8
|
+
|
|
9
|
+
Components
|
|
10
|
+
----------
|
|
11
|
+
|
|
12
|
+
There are two main components:
|
|
13
|
+
|
|
14
|
+
- **BraintrustPlugin**: Use this for both Temporal clients and workers. It's a convenience
|
|
15
|
+
wrapper that automatically configures the interceptor and sandbox settings.
|
|
16
|
+
|
|
17
|
+
- **BraintrustInterceptor**: The underlying interceptor. You can use this directly if you
|
|
18
|
+
need more control, but ``BraintrustPlugin`` is recommended for most use cases.
|
|
19
|
+
|
|
20
|
+
Worker Setup
|
|
21
|
+
------------
|
|
22
|
+
|
|
23
|
+
Use ``BraintrustPlugin`` when creating a worker::
|
|
24
|
+
|
|
25
|
+
import braintrust
|
|
26
|
+
from braintrust.contrib.temporal import BraintrustPlugin
|
|
27
|
+
from temporalio.client import Client
|
|
28
|
+
from temporalio.worker import Worker
|
|
29
|
+
|
|
30
|
+
braintrust.init_logger(project="my-project")
|
|
31
|
+
|
|
32
|
+
client = await Client.connect("localhost:7233")
|
|
33
|
+
|
|
34
|
+
worker = Worker(
|
|
35
|
+
client,
|
|
36
|
+
task_queue="my-queue",
|
|
37
|
+
workflows=[MyWorkflow],
|
|
38
|
+
activities=[my_activity],
|
|
39
|
+
plugins=[BraintrustPlugin()],
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
await worker.run()
|
|
43
|
+
|
|
44
|
+
Client Setup
|
|
45
|
+
------------
|
|
46
|
+
|
|
47
|
+
Use ``BraintrustPlugin`` when creating a client to propagate span context to workflows::
|
|
48
|
+
|
|
49
|
+
import braintrust
|
|
50
|
+
from braintrust.contrib.temporal import BraintrustPlugin
|
|
51
|
+
from temporalio.client import Client
|
|
52
|
+
|
|
53
|
+
braintrust.init_logger(project="my-project")
|
|
54
|
+
|
|
55
|
+
client = await Client.connect(
|
|
56
|
+
"localhost:7233",
|
|
57
|
+
plugins=[BraintrustPlugin()],
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
# Spans created around workflow calls will be linked as parents
|
|
61
|
+
with braintrust.start_span(name="my-operation") as span:
|
|
62
|
+
result = await client.execute_workflow(
|
|
63
|
+
MyWorkflow.run,
|
|
64
|
+
args,
|
|
65
|
+
id="workflow-id",
|
|
66
|
+
task_queue="my-queue",
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
What Gets Traced
|
|
70
|
+
----------------
|
|
71
|
+
|
|
72
|
+
The integration will automatically:
|
|
73
|
+
|
|
74
|
+
- Trace workflow executions
|
|
75
|
+
- Trace all activity executions
|
|
76
|
+
- Trace local activities
|
|
77
|
+
- Maintain parent-child relationships between client calls, workflows, and activities
|
|
78
|
+
- Handle child workflows
|
|
79
|
+
- Respect Temporal replay safety (no duplicate spans during replay)
|
|
80
|
+
"""
|
|
81
|
+
|
|
82
|
+
import dataclasses
|
|
83
|
+
from typing import Any, Dict, Mapping, Optional, Type
|
|
84
|
+
|
|
85
|
+
import braintrust
|
|
86
|
+
import temporalio.activity
|
|
87
|
+
import temporalio.api.common.v1
|
|
88
|
+
import temporalio.client
|
|
89
|
+
import temporalio.converter
|
|
90
|
+
import temporalio.worker
|
|
91
|
+
import temporalio.workflow
|
|
92
|
+
from temporalio.plugin import SimplePlugin
|
|
93
|
+
from temporalio.worker import WorkflowRunner
|
|
94
|
+
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
|
|
95
|
+
|
|
96
|
+
# Braintrust dynamically chooses its context implementation at runtime based on
|
|
97
|
+
# BRAINTRUST_OTEL_COMPAT environment variable. When first accessed, it reads
|
|
98
|
+
# os.environ which is restricted in the sandbox. Therefore if the first use
|
|
99
|
+
# is inside the sandbox, it will fail. So we eagerly reference it here to
|
|
100
|
+
# force initialization at import time (before sandbox evaluation).
|
|
101
|
+
try:
|
|
102
|
+
braintrust.current_span()
|
|
103
|
+
except Exception:
|
|
104
|
+
# It's okay if this fails (e.g., no logger initialized yet)
|
|
105
|
+
pass
|
|
106
|
+
|
|
107
|
+
# Store module-level reference to braintrust.current_span to avoid re-importing
|
|
108
|
+
# inside extern functions (which can trigger sandbox restrictions)
|
|
109
|
+
_current_span = braintrust.current_span
|
|
110
|
+
|
|
111
|
+
# Header key for passing span context between client, workflows, and activities
|
|
112
|
+
_HEADER_KEY = "_braintrust-span"
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor):
|
|
116
|
+
"""Braintrust interceptor for tracing Temporal workflows and activities.
|
|
117
|
+
|
|
118
|
+
This interceptor can be used with both Temporal clients and workers to automatically
|
|
119
|
+
trace workflow executions and activity runs. It maintains proper parent-child
|
|
120
|
+
relationships in the trace hierarchy and respects Temporal's replay safety requirements.
|
|
121
|
+
|
|
122
|
+
The interceptor:
|
|
123
|
+
- Creates spans for workflow executions (using sandbox_unrestricted)
|
|
124
|
+
- Captures activity execution as spans with metadata
|
|
125
|
+
- Propagates span context from client → workflow → activities
|
|
126
|
+
- Handles both regular activities and local activities
|
|
127
|
+
- Supports child workflows
|
|
128
|
+
- Logs errors from failed activities and workflows
|
|
129
|
+
- Ensures replay safety (no duplicate spans during workflow replay)
|
|
130
|
+
"""
|
|
131
|
+
|
|
132
|
+
def __init__(self, logger: Optional[Any] = None) -> None:
|
|
133
|
+
"""Initialize interceptor.
|
|
134
|
+
|
|
135
|
+
Args:
|
|
136
|
+
logger: Optional background logger for testing.
|
|
137
|
+
"""
|
|
138
|
+
self.payload_converter = temporalio.converter.PayloadConverter.default
|
|
139
|
+
self._bg_logger = logger
|
|
140
|
+
# Capture logger instance at init time for cross-thread use
|
|
141
|
+
if logger:
|
|
142
|
+
braintrust.logger._state._override_bg_logger.logger = logger
|
|
143
|
+
self._logger = braintrust.current_logger()
|
|
144
|
+
|
|
145
|
+
def _get_logger(self) -> Optional[Any]:
|
|
146
|
+
"""Get logger for creating spans.
|
|
147
|
+
|
|
148
|
+
Sets thread-local override if background logger provided (for testing),
|
|
149
|
+
then returns captured logger instance.
|
|
150
|
+
"""
|
|
151
|
+
if self._bg_logger:
|
|
152
|
+
braintrust.logger._state._override_bg_logger.logger = self._bg_logger
|
|
153
|
+
return self._logger
|
|
154
|
+
|
|
155
|
+
def intercept_client(
|
|
156
|
+
self, next: temporalio.client.OutboundInterceptor
|
|
157
|
+
) -> temporalio.client.OutboundInterceptor:
|
|
158
|
+
"""Intercept client calls to propagate span context to workflows."""
|
|
159
|
+
return _BraintrustClientOutboundInterceptor(next, self)
|
|
160
|
+
|
|
161
|
+
def intercept_activity(
|
|
162
|
+
self, next: temporalio.worker.ActivityInboundInterceptor
|
|
163
|
+
) -> temporalio.worker.ActivityInboundInterceptor:
|
|
164
|
+
"""Intercept activity executions to create activity spans."""
|
|
165
|
+
return _BraintrustActivityInboundInterceptor(next, self)
|
|
166
|
+
|
|
167
|
+
def workflow_interceptor_class(
|
|
168
|
+
self, input: temporalio.worker.WorkflowInterceptorClassInput
|
|
169
|
+
) -> Optional[Type["BraintrustWorkflowInboundInterceptor"]]:
|
|
170
|
+
"""Return workflow interceptor class to propagate context to activities."""
|
|
171
|
+
input.unsafe_extern_functions["__braintrust_get_logger"] = self._get_logger
|
|
172
|
+
return BraintrustWorkflowInboundInterceptor
|
|
173
|
+
|
|
174
|
+
def _span_context_to_headers(
|
|
175
|
+
self,
|
|
176
|
+
span_context: Dict[str, Any],
|
|
177
|
+
headers: Mapping[str, temporalio.api.common.v1.Payload],
|
|
178
|
+
) -> Mapping[str, temporalio.api.common.v1.Payload]:
|
|
179
|
+
"""Add span context to headers."""
|
|
180
|
+
if span_context:
|
|
181
|
+
payloads = self.payload_converter.to_payloads([span_context])
|
|
182
|
+
if payloads:
|
|
183
|
+
headers = {
|
|
184
|
+
**headers,
|
|
185
|
+
_HEADER_KEY: payloads[0],
|
|
186
|
+
}
|
|
187
|
+
return headers
|
|
188
|
+
|
|
189
|
+
def _span_context_from_headers(
|
|
190
|
+
self, headers: Mapping[str, temporalio.api.common.v1.Payload]
|
|
191
|
+
) -> Optional[Dict[str, Any]]:
|
|
192
|
+
"""Extract span context from headers."""
|
|
193
|
+
if _HEADER_KEY not in headers:
|
|
194
|
+
return None
|
|
195
|
+
header_payload = headers.get(_HEADER_KEY)
|
|
196
|
+
if not header_payload:
|
|
197
|
+
return None
|
|
198
|
+
payloads = self.payload_converter.from_payloads([header_payload])
|
|
199
|
+
if not payloads:
|
|
200
|
+
return None
|
|
201
|
+
return payloads[0] if payloads[0] else None
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
class _BraintrustClientOutboundInterceptor(temporalio.client.OutboundInterceptor):
|
|
205
|
+
"""Client interceptor that propagates span context to workflows."""
|
|
206
|
+
|
|
207
|
+
def __init__(
|
|
208
|
+
self, next: temporalio.client.OutboundInterceptor, root: BraintrustInterceptor
|
|
209
|
+
) -> None:
|
|
210
|
+
super().__init__(next)
|
|
211
|
+
self.root = root
|
|
212
|
+
|
|
213
|
+
async def start_workflow(
|
|
214
|
+
self, input: temporalio.client.StartWorkflowInput
|
|
215
|
+
) -> temporalio.client.WorkflowHandle[Any, Any]:
|
|
216
|
+
# Get current span context and add it to workflow headers
|
|
217
|
+
current_span = _current_span()
|
|
218
|
+
if current_span:
|
|
219
|
+
span_context = current_span.export()
|
|
220
|
+
input.headers = self.root._span_context_to_headers(span_context, input.headers)
|
|
221
|
+
|
|
222
|
+
return await super().start_workflow(input)
|
|
223
|
+
|
|
224
|
+
|
|
225
|
+
class _BraintrustActivityInboundInterceptor(temporalio.worker.ActivityInboundInterceptor):
|
|
226
|
+
"""Activity interceptor that creates spans for activity executions."""
|
|
227
|
+
|
|
228
|
+
def __init__(
|
|
229
|
+
self,
|
|
230
|
+
next: temporalio.worker.ActivityInboundInterceptor,
|
|
231
|
+
root: BraintrustInterceptor,
|
|
232
|
+
) -> None:
|
|
233
|
+
super().__init__(next)
|
|
234
|
+
self.root = root
|
|
235
|
+
|
|
236
|
+
async def execute_activity(
|
|
237
|
+
self, input: temporalio.worker.ExecuteActivityInput
|
|
238
|
+
) -> Any:
|
|
239
|
+
info = temporalio.activity.info()
|
|
240
|
+
|
|
241
|
+
# Extract parent span context from headers
|
|
242
|
+
parent_span_context = self.root._span_context_from_headers(input.headers)
|
|
243
|
+
|
|
244
|
+
logger = self.root._get_logger()
|
|
245
|
+
if not logger:
|
|
246
|
+
return await super().execute_activity(input)
|
|
247
|
+
|
|
248
|
+
# Create Braintrust span for activity execution, linked to workflow span
|
|
249
|
+
span = logger.start_span(
|
|
250
|
+
name=f"temporal.activity.{info.activity_type}",
|
|
251
|
+
type="task",
|
|
252
|
+
parent=parent_span_context or None,
|
|
253
|
+
metadata={
|
|
254
|
+
"temporal.activity_type": info.activity_type,
|
|
255
|
+
"temporal.activity_id": info.activity_id,
|
|
256
|
+
"temporal.workflow_id": info.workflow_id,
|
|
257
|
+
"temporal.workflow_run_id": info.workflow_run_id,
|
|
258
|
+
},
|
|
259
|
+
)
|
|
260
|
+
span.set_current()
|
|
261
|
+
|
|
262
|
+
try:
|
|
263
|
+
result = await super().execute_activity(input)
|
|
264
|
+
return result
|
|
265
|
+
except Exception as e:
|
|
266
|
+
span.log(error=str(e))
|
|
267
|
+
raise
|
|
268
|
+
finally:
|
|
269
|
+
span.unset_current()
|
|
270
|
+
span.end()
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
class BraintrustWorkflowInboundInterceptor(temporalio.worker.WorkflowInboundInterceptor):
|
|
274
|
+
"""Workflow interceptor that creates workflow spans and propagates context to activities.
|
|
275
|
+
|
|
276
|
+
This interceptor creates a span for the workflow execution using sandbox_unrestricted
|
|
277
|
+
to bypass Temporal's sandbox restrictions. The workflow span is the parent for all
|
|
278
|
+
activities and child workflows executed within it.
|
|
279
|
+
"""
|
|
280
|
+
|
|
281
|
+
def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None:
|
|
282
|
+
super().__init__(next)
|
|
283
|
+
self.payload_converter = temporalio.converter.PayloadConverter.default
|
|
284
|
+
self._parent_span_context: Optional[Dict[str, Any]] = None
|
|
285
|
+
|
|
286
|
+
def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
|
|
287
|
+
super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self))
|
|
288
|
+
|
|
289
|
+
async def execute_workflow(
|
|
290
|
+
self, input: temporalio.worker.ExecuteWorkflowInput
|
|
291
|
+
) -> Any:
|
|
292
|
+
# Extract parent span context from workflow headers (set by client)
|
|
293
|
+
parent_span_context = None
|
|
294
|
+
if _HEADER_KEY in input.headers:
|
|
295
|
+
header_payload = input.headers.get(_HEADER_KEY)
|
|
296
|
+
if header_payload:
|
|
297
|
+
payloads = self.payload_converter.from_payloads([header_payload])
|
|
298
|
+
if payloads:
|
|
299
|
+
parent_span_context = payloads[0]
|
|
300
|
+
|
|
301
|
+
# Store parent span context for activities (will be overwritten if we create a workflow span)
|
|
302
|
+
self._parent_span_context = parent_span_context
|
|
303
|
+
|
|
304
|
+
# Create a span for the workflow execution using sandbox_unrestricted
|
|
305
|
+
# to bypass the sandbox restrictions on logger state access
|
|
306
|
+
span = None
|
|
307
|
+
if not temporalio.workflow.unsafe.is_replaying():
|
|
308
|
+
with temporalio.workflow.unsafe.sandbox_unrestricted():
|
|
309
|
+
# Get logger via extern function (supports test logger parameter)
|
|
310
|
+
get_logger = temporalio.workflow.extern_functions()["__braintrust_get_logger"]
|
|
311
|
+
logger = get_logger()
|
|
312
|
+
|
|
313
|
+
if logger:
|
|
314
|
+
info = temporalio.workflow.info()
|
|
315
|
+
span = logger.start_span(
|
|
316
|
+
name=f"temporal.workflow.{info.workflow_type}",
|
|
317
|
+
type="task",
|
|
318
|
+
parent=parent_span_context or None,
|
|
319
|
+
metadata={
|
|
320
|
+
"temporal.workflow_type": info.workflow_type,
|
|
321
|
+
"temporal.workflow_id": info.workflow_id,
|
|
322
|
+
"temporal.run_id": info.run_id,
|
|
323
|
+
},
|
|
324
|
+
)
|
|
325
|
+
span.set_current()
|
|
326
|
+
|
|
327
|
+
# Update parent span context for activities
|
|
328
|
+
self._parent_span_context = span.export()
|
|
329
|
+
|
|
330
|
+
try:
|
|
331
|
+
result = await super().execute_workflow(input)
|
|
332
|
+
return result
|
|
333
|
+
except Exception as e:
|
|
334
|
+
if span:
|
|
335
|
+
with temporalio.workflow.unsafe.sandbox_unrestricted():
|
|
336
|
+
span.log(error=str(e))
|
|
337
|
+
raise
|
|
338
|
+
finally:
|
|
339
|
+
if span:
|
|
340
|
+
with temporalio.workflow.unsafe.sandbox_unrestricted():
|
|
341
|
+
span.unset_current()
|
|
342
|
+
span.end()
|
|
343
|
+
|
|
344
|
+
|
|
345
|
+
class _BraintrustWorkflowOutboundInterceptor(
|
|
346
|
+
temporalio.worker.WorkflowOutboundInterceptor
|
|
347
|
+
):
|
|
348
|
+
"""Outbound workflow interceptor that propagates span context to activities."""
|
|
349
|
+
|
|
350
|
+
def __init__(
|
|
351
|
+
self,
|
|
352
|
+
next: temporalio.worker.WorkflowOutboundInterceptor,
|
|
353
|
+
root: BraintrustWorkflowInboundInterceptor,
|
|
354
|
+
) -> None:
|
|
355
|
+
super().__init__(next)
|
|
356
|
+
self.root = root
|
|
357
|
+
|
|
358
|
+
def _add_span_context_to_headers(
|
|
359
|
+
self, headers: Mapping[str, temporalio.api.common.v1.Payload]
|
|
360
|
+
) -> Mapping[str, temporalio.api.common.v1.Payload]:
|
|
361
|
+
"""Add parent span context to headers if available.
|
|
362
|
+
|
|
363
|
+
Note: We always pass span context through headers, even during replay,
|
|
364
|
+
so activities can maintain proper parent-child relationships. The replay
|
|
365
|
+
safety is handled in the activity interceptor, which only creates spans
|
|
366
|
+
when the activity actually executes (not during replay).
|
|
367
|
+
"""
|
|
368
|
+
if self.root._parent_span_context:
|
|
369
|
+
payloads = self.root.payload_converter.to_payloads([self.root._parent_span_context])
|
|
370
|
+
if payloads:
|
|
371
|
+
return {**headers, _HEADER_KEY: payloads[0]}
|
|
372
|
+
return headers
|
|
373
|
+
|
|
374
|
+
def start_activity(
|
|
375
|
+
self, input: temporalio.worker.StartActivityInput
|
|
376
|
+
) -> temporalio.workflow.ActivityHandle:
|
|
377
|
+
input.headers = self._add_span_context_to_headers(input.headers)
|
|
378
|
+
return super().start_activity(input)
|
|
379
|
+
|
|
380
|
+
def start_local_activity(
|
|
381
|
+
self, input: temporalio.worker.StartLocalActivityInput
|
|
382
|
+
) -> temporalio.workflow.ActivityHandle:
|
|
383
|
+
input.headers = self._add_span_context_to_headers(input.headers)
|
|
384
|
+
return super().start_local_activity(input)
|
|
385
|
+
|
|
386
|
+
def start_child_workflow(
|
|
387
|
+
self, input: temporalio.worker.StartChildWorkflowInput
|
|
388
|
+
) -> temporalio.workflow.ChildWorkflowHandle:
|
|
389
|
+
input.headers = self._add_span_context_to_headers(input.headers)
|
|
390
|
+
return super().start_child_workflow(input)
|
|
391
|
+
|
|
392
|
+
|
|
393
|
+
def _modify_workflow_runner(existing: Optional[WorkflowRunner]) -> Optional[WorkflowRunner]:
|
|
394
|
+
"""Add braintrust to sandbox passthrough modules."""
|
|
395
|
+
if isinstance(existing, SandboxedWorkflowRunner):
|
|
396
|
+
new_restrictions = existing.restrictions.with_passthrough_modules("braintrust")
|
|
397
|
+
return dataclasses.replace(existing, restrictions=new_restrictions)
|
|
398
|
+
return existing
|
|
399
|
+
|
|
400
|
+
|
|
401
|
+
class BraintrustPlugin(SimplePlugin):
|
|
402
|
+
"""Braintrust plugin for Temporal that automatically configures tracing.
|
|
403
|
+
|
|
404
|
+
This plugin simplifies Braintrust integration with Temporal by:
|
|
405
|
+
- Automatically adding BraintrustInterceptor to the worker
|
|
406
|
+
- Configuring the sandbox to allow braintrust imports without unsafe.imports_passed_through()
|
|
407
|
+
|
|
408
|
+
Example usage:
|
|
409
|
+
from braintrust.contrib.temporal import BraintrustPlugin
|
|
410
|
+
from temporalio.worker import Worker
|
|
411
|
+
|
|
412
|
+
worker = Worker(
|
|
413
|
+
client,
|
|
414
|
+
task_queue="my-queue",
|
|
415
|
+
workflows=[MyWorkflow],
|
|
416
|
+
activities=[my_activity],
|
|
417
|
+
plugins=[BraintrustPlugin()],
|
|
418
|
+
)
|
|
419
|
+
|
|
420
|
+
Requires temporalio >= 1.19.0.
|
|
421
|
+
"""
|
|
422
|
+
|
|
423
|
+
def __init__(self, logger: Optional[Any] = None) -> None:
|
|
424
|
+
"""Initialize the Braintrust plugin.
|
|
425
|
+
|
|
426
|
+
Args:
|
|
427
|
+
logger: Optional background logger for testing.
|
|
428
|
+
"""
|
|
429
|
+
interceptor = BraintrustInterceptor(logger=logger)
|
|
430
|
+
super().__init__(
|
|
431
|
+
name="braintrust",
|
|
432
|
+
client_interceptors=[interceptor],
|
|
433
|
+
worker_interceptors=[interceptor],
|
|
434
|
+
workflow_runner=_modify_workflow_runner,
|
|
435
|
+
)
|
|
436
|
+
|
|
437
|
+
|
|
438
|
+
__all__ = ["BraintrustInterceptor", "BraintrustPlugin"]
|