Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 26 additions & 6 deletions py/src/braintrust/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,12 @@ class ParentSpanIds:
span_parents: list[str]


@dataclass(frozen=True)
class _SpanContextEntry:
span_object: Any
parent: "_SpanContextEntry | None"


class ContextManager(ABC):
"""Abstract base class for managing span context in Braintrust.

Expand Down Expand Up @@ -77,29 +83,43 @@ class BraintrustContextManager(ContextManager):
"""Braintrust-only context manager using contextvars when OTEL is not available."""

def __init__(self):
self._current_span: ContextVar[Any | None] = ContextVar("braintrust_current_span", default=None)
self._current_span: ContextVar[_SpanContextEntry | None] = ContextVar("braintrust_current_span", default=None)

def _get_active_entry(self) -> _SpanContextEntry | None:
entry = self._current_span.get()
# Async cleanup can end a span from a copied context without clearing
# the original context. Skip that stale span while preserving any active
# parent that preceded it.
while entry is not None and getattr(entry.span_object, "_logged_end_time", None) is not None:
entry = entry.parent
return entry

def get_current_span_info(self) -> SpanInfo | None:
"""Get information about the currently active span."""
current_span = self._current_span.get()
if not current_span:
entry = self._get_active_entry()
if entry is None:
return None

current_span = entry.span_object
# Return SpanInfo for BT spans
return SpanInfo(trace_id=current_span.root_span_id, span_id=current_span.span_id, span_object=current_span)

def get_parent_span_ids(self) -> ParentSpanIds | None:
"""Get parent information for creating a new Braintrust span."""
current_span = self._current_span.get()
if not current_span:
entry = self._get_active_entry()
if entry is None:
return None

current_span = entry.span_object
# If current span is a BT span, use it as parent
return ParentSpanIds(root_span_id=current_span.root_span_id, span_parents=[current_span.span_id])

def set_current_span(self, span_object: Any) -> Any:
"""Set the current active span."""
return self._current_span.set(span_object)
entry = None
if span_object is not None:
entry = _SpanContextEntry(span_object=span_object, parent=self._get_active_entry())
return self._current_span.set(entry)

def unset_current_span(self, context_token: Any = None) -> None:
"""Unset the current active span."""
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,302 @@
interactions:
- request:
body: '{"messages":[{"content":"What is 1 + 0?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '94'
content-type:
- application/json
host:
- api.openai.com
user-agent:
- AsyncOpenAI/Python 2.32.0
x-stainless-arch:
- arm64
x-stainless-async:
- async:asyncio
x-stainless-lang:
- python
x-stainless-os:
- MacOS
x-stainless-package-version:
- 2.32.0
x-stainless-raw-response:
- 'true'
x-stainless-retry-count:
- '0'
x-stainless-runtime:
- CPython
x-stainless-runtime-version:
- 3.14.6
method: POST
uri: https://api.openai.com/v1/chat/completions
response:
body:
string: "{\n \"id\": \"chatcmpl-E6yqVnsqAOvHpagpGLwLLoTLqyawV\",\n \"object\":
\"chat.completion\",\n \"created\": 1785332567,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
\"assistant\",\n \"content\": \"1 + 0 equals 1.\",\n \"refusal\":
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
\"default\",\n \"system_fingerprint\": \"fp_836208cd2a\"\n}\n"
headers:
access-control-expose-headers:
- X-Request-ID
- CF-Ray
- CF-Ray
alt-svc:
- h3=":443"; ma=86400
cf-cache-status:
- DYNAMIC
cf-ray:
- a22c8c483935ddfa-YYZ
connection:
- keep-alive
content-length:
- '821'
content-type:
- application/json
date:
- Wed, 29 Jul 2026 13:42:48 GMT
openai-processing-ms:
- '1015'
openai-version:
- '2020-10-01'
server:
- cloudflare
strict-transport-security:
- max-age=31536000; includeSubDomains; preload
transfer-encoding:
- chunked
x-content-type-options:
- nosniff
x-openai-proxy-wasm:
- v0.1
x-ratelimit-limit-requests:
- '30000'
x-ratelimit-limit-tokens:
- '150000000'
x-ratelimit-remaining-requests:
- '29999'
x-ratelimit-remaining-tokens:
- '149999995'
x-ratelimit-reset-requests:
- 2ms
x-ratelimit-reset-tokens:
- 0s
x-request-id:
- req_dbe641bded204be29da98e20199f6002
status:
code: 200
message: OK
- request:
body: '{"messages":[{"content":"What is 1 + 1?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '94'
content-type:
- application/json
host:
- api.openai.com
user-agent:
- AsyncOpenAI/Python 2.32.0
x-stainless-arch:
- arm64
x-stainless-async:
- async:asyncio
x-stainless-lang:
- python
x-stainless-os:
- MacOS
x-stainless-package-version:
- 2.32.0
x-stainless-raw-response:
- 'true'
x-stainless-retry-count:
- '0'
x-stainless-runtime:
- CPython
x-stainless-runtime-version:
- 3.14.6
method: POST
uri: https://api.openai.com/v1/chat/completions
response:
body:
string: "{\n \"id\": \"chatcmpl-E6yqWfKjPjvvbeg7Zl04J5l3vQDpc\",\n \"object\":
\"chat.completion\",\n \"created\": 1785332568,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
\"assistant\",\n \"content\": \"1 + 1 equals 2.\",\n \"refusal\":
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
\"default\",\n \"system_fingerprint\": \"fp_33a9b387c2\"\n}\n"
headers:
access-control-expose-headers:
- X-Request-ID
- CF-Ray
- CF-Ray
alt-svc:
- h3=":443"; ma=86400
cf-cache-status:
- DYNAMIC
cf-ray:
- a22c8c888b87ddfa-YYZ
connection:
- keep-alive
content-length:
- '821'
content-type:
- application/json
date:
- Wed, 29 Jul 2026 13:42:49 GMT
openai-processing-ms:
- '717'
openai-version:
- '2020-10-01'
server:
- cloudflare
strict-transport-security:
- max-age=31536000; includeSubDomains; preload
transfer-encoding:
- chunked
x-content-type-options:
- nosniff
x-openai-proxy-wasm:
- v0.1
x-ratelimit-limit-requests:
- '30000'
x-ratelimit-limit-tokens:
- '150000000'
x-ratelimit-remaining-requests:
- '29999'
x-ratelimit-remaining-tokens:
- '149999992'
x-ratelimit-reset-requests:
- 2ms
x-ratelimit-reset-tokens:
- 0s
x-request-id:
- req_63ce970e8407408dbfb043eb017efd93
status:
code: 200
message: OK
- request:
body: '{"messages":[{"content":"What is 1 + 2?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '94'
content-type:
- application/json
host:
- api.openai.com
user-agent:
- AsyncOpenAI/Python 2.32.0
x-stainless-arch:
- arm64
x-stainless-async:
- async:asyncio
x-stainless-lang:
- python
x-stainless-os:
- MacOS
x-stainless-package-version:
- 2.32.0
x-stainless-raw-response:
- 'true'
x-stainless-retry-count:
- '0'
x-stainless-runtime:
- CPython
x-stainless-runtime-version:
- 3.14.6
method: POST
uri: https://api.openai.com/v1/chat/completions
response:
body:
string: "{\n \"id\": \"chatcmpl-E6yqXMMW5UVWCocAJ6F4OlSRWpmuX\",\n \"object\":
\"chat.completion\",\n \"created\": 1785332569,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
\"assistant\",\n \"content\": \"1 + 2 equals 3.\",\n \"refusal\":
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
\"default\",\n \"system_fingerprint\": \"fp_ab0a2ab924\"\n}\n"
headers:
access-control-expose-headers:
- X-Request-ID
- CF-Ray
- CF-Ray
alt-svc:
- h3=":443"; ma=86400
cf-cache-status:
- DYNAMIC
cf-ray:
- a22c8c8e29baddfa-YYZ
connection:
- keep-alive
content-length:
- '821'
content-type:
- application/json
date:
- Wed, 29 Jul 2026 13:42:50 GMT
openai-processing-ms:
- '646'
openai-version:
- '2020-10-01'
server:
- cloudflare
strict-transport-security:
- max-age=31536000; includeSubDomains; preload
transfer-encoding:
- chunked
x-content-type-options:
- nosniff
x-openai-proxy-wasm:
- v0.1
x-ratelimit-limit-requests:
- '30000'
x-ratelimit-limit-tokens:
- '150000000'
x-ratelimit-remaining-requests:
- '29999'
x-ratelimit-remaining-tokens:
- '149999995'
x-ratelimit-reset-requests:
- 2ms
x-ratelimit-reset-tokens:
- 0s
x-request-id:
- req_1f3ebdbc09074c9a84a3b4a4001cf9ad
status:
code: 200
message: OK
version: 1
Loading