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
4 changes: 3 additions & 1 deletion py/src/braintrust/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -3200,6 +3200,7 @@ def __init__(
pinned_version: None | int | str = None,
mutate_record: Callable[[TMapping], TMapping] | None = None,
_internal_btql: dict[str, Any] | None = None,
_internal_brainstore_realtime: bool = True,
):
self.object_type = object_type

Expand All @@ -3215,6 +3216,7 @@ def __init__(

self._fetched_data: list[TMapping] | None = None
self._internal_btql = _internal_btql
self._internal_brainstore_realtime = _internal_brainstore_realtime

def fetch(self, batch_size: int | None = None) -> Iterator[TMapping]:
"""
Expand Down Expand Up @@ -3282,7 +3284,7 @@ def _refetch(self, batch_size: int | None = None) -> list[TMapping]:
**(self._internal_btql or {}),
},
"use_columnstore": False,
"brainstore_realtime": True,
"brainstore_realtime": self._internal_brainstore_realtime,
"query_source": f"py_sdk_object_fetcher_{self.object_type}",
**({"version": self._pinned_version} if self._pinned_version is not None else {}),
},
Expand Down
70 changes: 68 additions & 2 deletions py/src/braintrust/test_trace.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Tests for Trace functionality."""

import pytest
from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData
from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData, SpanFetcher


# Helper to create mock spans
Expand Down Expand Up @@ -328,19 +328,85 @@ async def fetch_fn(span_type):
assert call_args[0] is None or call_args[0] == []
assert len(result) == 1

@pytest.mark.parametrize(
("brainstore_realtime", "expected"),
[
(None, True),
(False, False),
],
)
def test_span_fetcher_threads_realtime_setting(self, brainstore_realtime, expected):
calls = []
state = _DummyState(calls)
kwargs = dict(
object_type="project_logs",
object_id="project-1",
root_span_id="root-1",
state=state,
)
if brainstore_realtime is not None:
kwargs["brainstore_realtime"] = brainstore_realtime
fetcher = SpanFetcher(**kwargs)

assert list(fetcher.fetch()) == []
assert calls[0]["json"]["brainstore_realtime"] is expected

@pytest.mark.asyncio
async def test_cached_span_fetcher_threads_realtime_setting(self):
calls = []
state = _DummyState(calls)

async def get_state():
return state

fetcher = CachedSpanFetcher(
object_type="project_logs",
object_id="project-1",
root_span_id="root-1",
get_state=get_state,
brainstore_realtime=False,
)

assert await fetcher.get_spans() == []
assert calls[0]["json"]["brainstore_realtime"] is False


class _DummySpanCache:
def get_by_root_span_id(self, root_span_id: str):
return None


class _DummyState:
def __init__(self):
def __init__(self, api_calls=None):
self.span_cache = _DummySpanCache()
self.api_calls = api_calls

def login(self):
return None

def api_conn(self):
return _DummyApiConn(self.api_calls)


class _DummyResponse:
text = ""

def raise_for_status(self):
return None

def json(self):
return {"data": []}


class _DummyApiConn:
def __init__(self, calls):
self.calls = calls

def post(self, path, *args, **kwargs):
if self.calls is not None:
self.calls.append({"path": path, "args": args, **kwargs})
return _DummyResponse()


class TestLocalTraceGetThread:
@pytest.mark.asyncio
Expand Down
4 changes: 4 additions & 0 deletions py/src/braintrust/trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,13 +75,15 @@ def __init__(
state: BraintrustState,
span_type_filter: list[str] | None = None,
include_scorers: bool = False,
brainstore_realtime: bool = True,
):
# Build the filter expression for root_span_id and optionally span_attributes.type
filter_expr = self._build_filter(root_span_id, span_type_filter, include_scorers)

super().__init__(
object_type=object_type,
_internal_btql={"filter": filter_expr},
_internal_brainstore_realtime=brainstore_realtime,
)
self._object_id = object_id
self._state = state
Expand Down Expand Up @@ -171,6 +173,7 @@ def __init__(
root_span_id: str | None = None,
get_state: Callable[[], Awaitable[BraintrustState]] | None = None,
fetch_fn: SpanFetchFn | None = None,
brainstore_realtime: bool = True,
):
self._span_cache: dict[str, list[SpanData]] = {}
self._all_fetched = False
Expand Down Expand Up @@ -204,6 +207,7 @@ async def _fetch_fn(
state=state,
span_type_filter=span_type,
include_scorers=include_scorers,
brainstore_realtime=brainstore_realtime,
)
rows = list(fetcher.fetch())
return [
Expand Down