From 0ddb869f54bb991abc3acbd182c6f0bf761bea22 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 1 Sep 2026 13:41:55 -0400 Subject: [PATCH] feat(trace): add advanced span filters Add nested `filters` to `Trace.get_spans()` so callers can select spans by type, name, error state, tags, metadata, and duration with matching local and BTQL semantics. ```python spans = await trace.get_spans( filters={ "span_type": ["tool"], "name": ["search", "lookup"], "has_error": True, "metadata": {"model": "gpt-5"}, "duration": {"min": 0.5, "max": 10}, } ) ``` Preserve fields needed to evaluate these filters in the local span cache. Keep existing full-trace and span-type caching, but execute advanced remote filters as fresh BTQL requests so partial or empty results do not become stale snapshots: ```text local buffered spans -> in-memory filtering complete trace cache -> in-memory filtering advanced remote filter -> fresh BTQL request ``` Keep `span_type=[...]` compatible on the public `Trace` API but emit a `DeprecationWarning` directing callers to `filters={"span_type": [...]}`. Add runtime and type coverage for validation and equivalent local and server filtering. --- py/src/braintrust/logger.py | 3 + py/src/braintrust/span_cache.py | 49 +-- py/src/braintrust/test_span_cache.py | 7 + py/src/braintrust/test_trace.py | 485 +++++++++++++--------- py/src/braintrust/trace.py | 581 +++++++++++++++++++++------ 5 files changed, 808 insertions(+), 317 deletions(-) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index ee4ba6779..d1752d4b4 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -4627,6 +4627,9 @@ def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[ metadata=serializable_partial_record.get("metadata"), span_parents=self.span_parents, span_attributes=serializable_partial_record.get("span_attributes"), + error=serializable_partial_record.get("error"), + metrics=serializable_partial_record.get("metrics"), + tags=serializable_partial_record.get("tags"), ) self.state.span_cache.queue_write(self.root_span_id, self.span_id, cached_span) diff --git a/py/src/braintrust/span_cache.py b/py/src/braintrust/span_cache.py index ee9266142..734804415 100644 --- a/py/src/braintrust/span_cache.py +++ b/py/src/braintrust/span_cache.py @@ -14,7 +14,7 @@ from typing import Any from braintrust.types import Metadata -from braintrust.util import merge_dicts +from braintrust.util import clean_nones, merge_dicts # Global registry of active span caches for process exit cleanup @@ -23,7 +23,12 @@ class CachedSpan: - """Cached span data structure.""" + """A span held in the local cache, before it has been flushed to the server. + + Carries the subset of span fields that scorers can filter on, so that a trace can be + queried without a round-trip. Fields the server has but this does not are simply not + filterable locally. + """ def __init__( self, @@ -33,6 +38,9 @@ def __init__( metadata: Metadata | None = None, span_parents: list[str] | None = None, span_attributes: dict[str, Any] | None = None, + error: Any | None = None, + metrics: dict[str, Any] | None = None, + tags: list[str] | None = None, ): self.span_id = span_id self.input = input @@ -40,33 +48,26 @@ def __init__( self.metadata = metadata self.span_parents = span_parents self.span_attributes = span_attributes + self.error = error + self.metrics = metrics + self.tags = tags def to_dict(self) -> dict[str, Any]: - """Convert to dictionary for serialization.""" - result = {"span_id": self.span_id} - if self.input is not None: - result["input"] = self.input - if self.output is not None: - result["output"] = self.output - if self.metadata is not None: - result["metadata"] = self.metadata - if self.span_parents is not None: - result["span_parents"] = self.span_parents - if self.span_attributes is not None: - result["span_attributes"] = self.span_attributes - return result + """Return the span's set fields, dropping those left as None. + + Unset fields are omitted rather than written as null to keep the on-disk record + small; span_id is always present, so it survives the stripping. + """ + return clean_nones(self.__dict__) @classmethod def from_dict(cls, data: dict[str, Any]) -> "CachedSpan": - """Create from dictionary.""" - return cls( - span_id=data["span_id"], - input=data.get("input"), - output=data.get("output"), - metadata=data.get("metadata"), - span_parents=data.get("span_parents"), - span_attributes=data.get("span_attributes"), - ) + """Rebuild a span from a record produced by to_dict(). + + The cache file is written and read by one process, so `data` always has exactly the + fields this class defines and can be passed straight through. + """ + return cls(**data) class DiskSpanRecord: diff --git a/py/src/braintrust/test_span_cache.py b/py/src/braintrust/test_span_cache.py index 9b250d445..767d169ae 100644 --- a/py/src/braintrust/test_span_cache.py +++ b/py/src/braintrust/test_span_cache.py @@ -13,6 +13,9 @@ def test_span_cache_write_and_read(): span_id="span-1", input={"text": "hello"}, output={"response": "world"}, + error={"message": "retryable"}, + metrics={"start": 1, "end": 3}, + tags=["production"], ) span2 = CachedSpan( span_id="span-2", @@ -30,6 +33,10 @@ def test_span_cache_write_and_read(): span_ids = {s.span_id for s in spans} assert "span-1" in span_ids assert "span-2" in span_ids + stored_span1 = next(span for span in spans if span.span_id == "span-1") + assert stored_span1.error == {"message": "retryable"} + assert stored_span1.metrics == {"start": 1, "end": 3} + assert stored_span1.tags == ["production"] cache.stop() cache.dispose() diff --git a/py/src/braintrust/test_trace.py b/py/src/braintrust/test_trace.py index 3a572f68e..ac2c99e11 100644 --- a/py/src/braintrust/test_trace.py +++ b/py/src/braintrust/test_trace.py @@ -1,20 +1,39 @@ """Tests for Trace functionality.""" import pytest -from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData, SpanFetcher +from braintrust.span_cache import CachedSpan +from braintrust.trace import ( + _FILTER_SPECS, + CachedSpanFetcher, + LocalTrace, + SpanData, + SpanFetcher, + SpanFilters, + _matches_span_filters, +) # Helper to create mock spans -def make_span(span_id: str, span_type: str, **extra) -> SpanData: +def make_span(span_id: str, span_type: str, *, name: str | None = None, **extra) -> SpanData: + span_attributes = {"type": span_type} + if name is not None: + span_attributes["name"] = name return SpanData( span_id=span_id, input={"text": f"input-{span_id}"}, output={"text": f"output-{span_id}"}, - span_attributes={"type": span_type}, + span_attributes=span_attributes, **extra, ) +def test_every_span_filter_field_has_a_spec(): + """SpanFilters is the public contract; _FILTER_SPECS is what implements it.""" + assert set(_FILTER_SPECS) == set(SpanFilters.__annotations__) + # Registration order fixes BTQL child order, so it must track the declaration order. + assert list(_FILTER_SPECS) == list(SpanFilters.__annotations__) + + class TestCachedSpanFetcher: """Test CachedSpanFetcher caching behavior.""" @@ -29,7 +48,7 @@ async def test_fetch_all_spans_without_filter(self): call_count = 0 - async def fetch_fn(span_type): + async def fetch_fn(filters): nonlocal call_count call_count += 1 return mock_spans @@ -50,13 +69,14 @@ async def test_fetch_all_after_typed_fetch_has_no_duplicates(self): make_span("llm-2", "llm"), ] - async def fetch_fn(span_type): + async def fetch_fn(filters): + span_type = filters.get("span_type") if span_type: return [s for s in all_spans if s.span_attributes["type"] in span_type] return all_spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - await fetcher.get_spans(["llm"]) + await fetcher.get_spans(filters={"span_type": ["llm"]}) result = await fetcher.get_spans() span_ids = [s.span_id for s in result] @@ -78,8 +98,8 @@ async def test_fetch_preserves_span_result_fields(self): ) ] - async def fetch_fn(span_type): - del span_type + async def fetch_fn(filters): + del filters return mock_spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) @@ -99,167 +119,74 @@ async def test_fetch_specific_span_types(self): call_count = 0 - async def fetch_fn(span_type): - nonlocal call_count - call_count += 1 - assert span_type == ["llm"] - return llm_spans - - fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - result = await fetcher.get_spans(span_type=["llm"]) - - assert call_count == 1 - assert len(result) == 2 - - @pytest.mark.asyncio - async def test_return_cached_spans_after_fetching_all(self): - """Test that cached spans are returned without re-fetching after fetching all.""" - mock_spans = [ - make_span("span-1", "llm"), - make_span("span-2", "function"), - ] - - call_count = 0 - - async def fetch_fn(span_type): - nonlocal call_count - call_count += 1 - return mock_spans - - fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - - # First call - fetches - await fetcher.get_spans() - assert call_count == 1 - - # Second call - should use cache - result = await fetcher.get_spans() - assert call_count == 1 # Still 1 - assert len(result) == 2 - - @pytest.mark.asyncio - async def test_return_cached_spans_for_previously_fetched_types(self): - """Test that previously fetched types are returned from cache.""" - llm_spans = [make_span("span-1", "llm"), make_span("span-2", "llm")] - - call_count = 0 - - async def fetch_fn(span_type): + async def fetch_fn(filters): nonlocal call_count call_count += 1 + assert filters == {"span_type": ["llm"]} return llm_spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + result = await fetcher.get_spans(filters={"span_type": ["llm"]}) - # First call - fetches llm spans - await fetcher.get_spans(span_type=["llm"]) assert call_count == 1 - - # Second call for same type - should use cache - result = await fetcher.get_spans(span_type=["llm"]) - assert call_count == 1 # Still 1 assert len(result) == 2 + @pytest.mark.parametrize( + ("span_type", "expected_ids"), + [ + (None, ["span-1", "span-2", "span-3", "span-4"]), + (["llm"], ["span-1", "span-4"]), + (["llm", "tool"], ["span-1", "span-3", "span-4"]), + (["nonexistent"], []), + ], + ) @pytest.mark.asyncio - async def test_only_fetch_missing_span_types(self): - """Test that only missing span types are fetched.""" - llm_spans = [make_span("span-1", "llm")] - function_spans = [make_span("span-2", "function")] + async def test_full_cache_answers_any_span_type_query(self, span_type, expected_ids): + """One unfiltered fetch makes the cache authoritative for every span type. - call_count = 0 - - async def fetch_fn(span_type): - nonlocal call_count - call_count += 1 - if span_type == ["llm"]: - return llm_spans - elif span_type == ["function"]: - return function_spans - return [] - - fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - - # First call - fetches llm spans - await fetcher.get_spans(span_type=["llm"]) - assert call_count == 1 - - # Second call for both types - should only fetch function - result = await fetcher.get_spans(span_type=["llm", "function"]) - assert call_count == 2 - assert len(result) == 2 - - @pytest.mark.asyncio - async def test_no_refetch_after_fetching_all_spans(self): - """Test that no re-fetching occurs after fetching all spans.""" + Including types that turn out to be absent: an empty result is a real answer here, + not a cache miss to be retried against the server. + """ all_spans = [ make_span("span-1", "llm"), make_span("span-2", "function"), make_span("span-3", "tool"), + make_span("span-4", "llm"), ] - call_count = 0 - async def fetch_fn(span_type): + async def fetch_fn(filters): nonlocal call_count call_count += 1 return all_spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - - # Fetch all spans await fetcher.get_spans() - assert call_count == 1 - # Subsequent filtered calls should use cache - llm_result = await fetcher.get_spans(span_type=["llm"]) - assert call_count == 1 # Still 1 - assert len(llm_result) == 1 - assert llm_result[0].span_id == "span-1" + result = await fetcher.get_spans(filters={"span_type": span_type} if span_type else None) - function_result = await fetcher.get_spans(span_type=["function"]) - assert call_count == 1 # Still 1 - assert len(function_result) == 1 - assert function_result[0].span_id == "span-2" - - @pytest.mark.asyncio - async def test_filter_by_multiple_span_types_from_cache(self): - """Test filtering by multiple span types from cache.""" - all_spans = [ - make_span("span-1", "llm"), - make_span("span-2", "function"), - make_span("span-3", "tool"), - make_span("span-4", "llm"), - ] - - async def fetch_fn(span_type): - return all_spans - - fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - - # Fetch all first - await fetcher.get_spans() - - # Filter for llm and tool - result = await fetcher.get_spans(span_type=["llm", "tool"]) - assert len(result) == 3 - assert {s.span_id for s in result} == {"span-1", "span-3", "span-4"} + assert call_count == 1 + assert sorted(span.span_id for span in result) == expected_ids @pytest.mark.asyncio - async def test_return_empty_for_nonexistent_span_type(self): - """Test that empty array is returned for non-existent span type.""" - all_spans = [make_span("span-1", "llm")] + async def test_partial_cache_fetches_only_missing_types(self): + """A type already in the cache is never re-requested, only the types missing from it.""" + by_type = {"llm": [make_span("span-1", "llm")], "function": [make_span("span-2", "function")]} + requested = [] - async def fetch_fn(span_type): - return all_spans + async def fetch_fn(filters): + requested.append(filters["span_type"]) + return [span for t in filters["span_type"] for span in by_type.get(t, [])] fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - # Fetch all first - await fetcher.get_spans() + assert [s.span_id for s in await fetcher.get_spans(filters={"span_type": ["llm"]})] == ["span-1"] + assert [s.span_id for s in await fetcher.get_spans(filters={"span_type": ["llm"]})] == ["span-1"] + result = await fetcher.get_spans(filters={"span_type": ["llm", "function"]}) - # Query for non-existent type - result = await fetcher.get_spans(span_type=["nonexistent"]) - assert len(result) == 0 + assert sorted(span.span_id for span in result) == ["span-1", "span-2"] + # The second call was served from cache; the third asked only for what it lacked. + assert requested == [["llm"], ["function"]] @pytest.mark.asyncio async def test_handle_spans_with_no_type(self): @@ -270,7 +197,7 @@ async def test_handle_spans_with_no_type(self): SpanData(span_id="span-3", input={}), # No span_attributes ] - async def fetch_fn(span_type): + async def fetch_fn(filters): return spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) @@ -280,75 +207,178 @@ async def fetch_fn(span_type): assert len(result) == 3 # Spans without type go into "" bucket - no_type_result = await fetcher.get_spans(span_type=[""]) + no_type_result = await fetcher.get_spans(filters={"span_type": [""]}) assert len(no_type_result) == 2 + @pytest.mark.parametrize("filters", [None, {"span_type": ["llm"]}]) @pytest.mark.asyncio - async def test_empty_then_populated_refetches(self): - """Test that empty results don't permanently cache, allowing re-fetch when data becomes available.""" + async def test_empty_results_are_not_cached(self, filters): + """An empty fetch caches nothing, so spans logged later are still picked up. + + The cache records which types it holds by the spans it saw, so a fetch that returned + nothing leaves no trace and the next call goes back to the server. + """ call_count = 0 - spans = [make_span("span-1", "llm"), make_span("span-2", "function")] - async def fetch_fn(span_type): + async def fetch_fn(_filters): nonlocal call_count call_count += 1 - if call_count == 1: - return [] - return spans + return [] if call_count == 1 else [make_span("span-1", "llm")] fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - # First call returns empty - result1 = await fetcher.get_spans() - assert len(result1) == 0 - assert call_count == 1 - - # Second call should re-fetch since first was empty - result2 = await fetcher.get_spans() + assert await fetcher.get_spans(filters=filters) == [] + assert [span.span_id for span in await fetcher.get_spans(filters=filters)] == ["span-1"] assert call_count == 2 - assert len(result2) == 2 - assert {s.span_id for s in result2} == {"span-1", "span-2"} + + def test_span_fetcher_builds_advanced_filter(self): + calls = [] + state = _DummyState(calls) + fetcher = SpanFetcher( + object_type="project_logs", + object_id="project-1", + root_span_id="root-1", + state=state, + filters={ + "span_type": ["tool"], + "name": ["search", "lookup"], + "has_error": False, + "metadata": {"model": "gpt-5", "optional": None, "request": {"region": "us-east-1"}}, + "duration": {"min": 0.5, "max": 10}, + }, + ) + + assert list(fetcher.fetch()) == [] + + def comparison(op, path, value): + return { + "op": op, + "left": {"op": "ident", "name": path}, + "right": {"op": "literal", "value": value}, + } + + def isnull(path): + return {"op": "isnull", "expr": {"op": "ident", "name": path}} + + duration = { + "op": "sub", + "left": {"op": "ident", "name": ["metrics", "end"]}, + "right": {"op": "ident", "name": ["metrics", "start"]}, + } + purpose = ["span_attributes", "purpose"] + + # Asserting the whole expression pins the child order too, so adding a filter + # cannot silently reorder or drop one. + assert calls[0]["json"]["query"]["filter"] == { + "op": "and", + "children": [ + comparison("eq", ["root_span_id"], "root-1"), + {"op": "or", "children": [isnull(purpose), comparison("ne", purpose, "scorer")]}, + comparison("in", ["span_attributes", "type"], ["tool"]), + comparison("in", ["span_attributes", "name"], ["search", "lookup"]), + isnull(["error"]), + comparison("eq", ["metadata", "model"], "gpt-5"), + isnull(["metadata", "optional"]), + comparison("eq", ["metadata", "request", "region"], "us-east-1"), + {"op": "ge", "left": duration, "right": {"op": "literal", "value": 0.5}}, + {"op": "le", "left": duration, "right": {"op": "literal", "value": 10}}, + ], + } @pytest.mark.asyncio - async def test_empty_results_with_type_filter(self): - """Test that type-filtered fetches handle empty results correctly.""" + async def test_advanced_filters_use_full_cache_when_available(self): + spans = [ + make_span( + "matching", + "tool", + name="search", + error={"message": "boom"}, + metadata={"model": "gpt-5", "request": {"region": "us-east-1", "id": 1}}, + metrics={"start": 1, "end": 4}, + ), + make_span( + "too-fast", + "tool", + name="search", + error={"message": "boom"}, + metadata={"model": "gpt-5"}, + metrics={"start": 1, "end": 1.1}, + ), + make_span("successful", "tool", name="search"), + ] call_count = 0 - async def fetch_fn(span_type): + async def fetch_fn(filters): nonlocal call_count + del filters call_count += 1 - if call_count == 1: - return [] - return [make_span("span-1", "llm")] + return spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + await fetcher.get_spans() + result = await fetcher.get_spans( + filters={ + "span_type": ["tool"], + "name": ["search"], + "has_error": True, + "metadata": {"request": {"region": "us-east-1"}}, + "duration": {"min": 2}, + } + ) - # First call with type filter returns empty - result1 = await fetcher.get_spans(span_type=["llm"]) - assert len(result1) == 0 - - # Second call with same type should re-fetch since type wasn't cached with results - result2 = await fetcher.get_spans(span_type=["llm"]) - assert call_count == 2 - assert len(result2) == 1 + assert call_count == 1 + assert [span.span_id for span in result] == ["matching"] @pytest.mark.asyncio - async def test_handle_empty_span_type_array(self): - """Test that empty spanType array is handled same as undefined.""" - mock_spans = [make_span("span-1", "llm")] + async def test_advanced_filters_are_pushed_down_and_never_cached(self): + """Filters the cache cannot reason about go to the fetcher whole, every time. - call_args = [] + The cache is partitioned by span type alone, so it cannot tell whether it holds + every span matching some other field. Rather than guess, these queries are pushed + down in full and their results are used once and discarded. + """ + spans = [ + make_span("errored", "tool", name="search", error={"message": "boom"}), + make_span("successful", "tool", name="search"), + ] + received = [] - async def fetch_fn(span_type): - call_args.append(span_type) - return mock_spans + async def fetch_fn(filters): + received.append(filters) + return [span for span in spans if _matches_span_filters(span, filters)] fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + filters = {"span_type": ["tool"], "has_error": True} + + first = await fetcher.get_spans(filters=filters) + second = await fetcher.get_spans(filters=filters) + + # Handed down whole, returned unchanged (no second, client-side filtering pass), + # and re-fetched rather than served from the first call's results. + assert received == [filters, filters] + assert [span.span_id for span in first] == ["errored"] + assert [span.span_id for span in second] == ["errored"] - result = await fetcher.get_spans(span_type=[]) + @pytest.mark.parametrize( + ("filters", "message"), + [ + ({"span_type": []}, "span_type"), + ({"name": []}, "name"), + ({"duration": {"min": -1}}, "duration.min"), + ({"duration": {"min": 2, "max": 1}}, "duration.min"), + # A field that is present must constrain something, at any nesting depth. + ({"metadata": {}}, "filters.metadata must not be empty"), + ({"metadata": {"a": {}}}, "filters.metadata.a must not be empty"), + ({"metadata": {"a": {"b": {}}}}, "filters.metadata.a.b must not be empty"), + ({"duration": {}}, "filters.duration must specify at least one of: min, max"), + ], + ) + @pytest.mark.asyncio + async def test_rejects_invalid_advanced_filters(self, filters, message): + fetcher = CachedSpanFetcher(fetch_fn=lambda filters: None) - assert call_args[0] is None or call_args[0] == [] - assert len(result) == 1 + with pytest.raises(ValueError, match=message): + await fetcher.get_spans(filters=filters) @pytest.mark.parametrize( ("brainstore_realtime", "expected"), @@ -393,14 +423,107 @@ async def get_state(): assert calls[0]["json"]["brainstore_realtime"] is False +class TestLocalTraceGetSpans: + @pytest.mark.asyncio + async def test_applies_advanced_filters_to_local_spans(self): + def cached_span(span_id, name, **attributes): + return CachedSpan( + span_id=span_id, + input={"query": "weather"}, + output={"result": "sunny"}, + error={"message": "retryable"}, + metrics={"start": 1, "end": 4}, + metadata={"request": {"region": "us-east-1", "id": 1}}, + span_parents=[], + span_attributes={"type": "tool", "name": name, **attributes}, + ) + + spans = [ + cached_span("matching", "search"), + cached_span("wrong-name", "lookup"), + cached_span("scorer", "search", purpose="scorer"), + ] + trace = LocalTrace( + object_type="project_logs", + object_id="project-1", + root_span_id="root-1", + ensure_spans_flushed=None, + state=_DummyState(spans=spans), + ) + + result = await trace.get_spans( + filters={ + "span_type": ["tool"], + "name": ["search"], + "has_error": True, + "metadata": {"request": {"region": "us-east-1"}}, + "duration": {"min": 2, "max": 5}, + } + ) + + assert [span.span_id for span in result] == ["matching"] + + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + legacy_result = await trace.get_spans(span_type=["tool"]) + assert [span.span_id for span in legacy_result] == ["matching", "wrong-name"] + + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + with pytest.raises(ValueError, match="span_type"): + await trace.get_spans(["llm"], filters={"span_type": ["tool"]}) + + @pytest.mark.asyncio + async def test_empty_filters_object_is_not_a_filter(self): + """filters={} constrains nothing, unlike a field that is present but empty.""" + spans = [CachedSpan(span_id="tool-span", span_attributes={"type": "tool"})] + trace = LocalTrace( + object_type="project_logs", + object_id="project-1", + root_span_id="root-1", + ensure_spans_flushed=None, + state=_DummyState(spans=spans), + ) + + assert [span.span_id for span in await trace.get_spans(filters={})] == ["tool-span"] + assert [span.span_id for span in await trace.get_spans()] == ["tool-span"] + + @pytest.mark.asyncio + async def test_empty_span_type_only_bypasses_filtering_on_the_deprecated_argument(self): + """span_type=[] keeps its legacy "no filter" meaning; filters={"span_type": []} does not.""" + spans = [ + CachedSpan(span_id="tool-span", span_attributes={"type": "tool"}), + CachedSpan(span_id="llm-span", span_attributes={"type": "llm"}), + ] + trace = LocalTrace( + object_type="project_logs", + object_id="project-1", + root_span_id="root-1", + ensure_spans_flushed=None, + state=_DummyState(spans=spans), + ) + + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + unfiltered = await trace.get_spans(span_type=[]) + assert [span.span_id for span in unfiltered] == ["tool-span", "llm-span"] + + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + filtered = await trace.get_spans(span_type=["tool"]) + assert [span.span_id for span in filtered] == ["tool-span"] + + with pytest.raises(ValueError, match="span_type"): + await trace.get_spans(filters={"span_type": []}) + + class _DummySpanCache: + def __init__(self, spans=None): + self.spans = spans + def get_by_root_span_id(self, root_span_id: str): - return None + return self.spans class _DummyState: - def __init__(self, api_calls=None): - self.span_cache = _DummySpanCache() + def __init__(self, api_calls=None, spans=None): + self.span_cache = _DummySpanCache(spans) self.api_calls = api_calls def login(self): diff --git a/py/src/braintrust/trace.py b/py/src/braintrust/trace.py index b72b8d621..7de9aab93 100644 --- a/py/src/braintrust/trace.py +++ b/py/src/braintrust/trace.py @@ -6,16 +6,389 @@ """ import asyncio -from collections.abc import Awaitable, Callable -from typing import Any, Protocol, TypedDict +import math +import warnings +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Mapping +from typing import Any, ClassVar, Protocol, TypedDict, cast from braintrust.functions.invoke import invoke from braintrust.logger import BraintrustState, ObjectFetcher from braintrust.types import Metadata +from braintrust.util import clean_nones, is_numeric + + +class SpanDurationFilter(TypedDict, total=False): + """Inclusive duration bounds, in seconds.""" + + min: float + """Minimum value of metrics.end - metrics.start.""" + max: float + """Maximum value of metrics.end - metrics.start.""" + + +class SpanFilters(TypedDict, total=False): + """Filters supported by Trace.get_spans(). Different fields combine with AND. + + Every field provided must constrain something: empty lists and empty objects are + rejected rather than ignored, so a filter that comes out empty fails loudly instead + of silently matching every span. Omit a field to leave it unfiltered. + """ + + span_type: list[str] + """Match spans whose span_attributes.type equals any of these.""" + name: list[str] + """Match spans whose span_attributes.name equals any of these.""" + has_error: bool + """True to keep only spans that recorded an error, False to keep only those that did not.""" + metadata: dict[str, Any] + """Match the named metadata keys, at any depth, leaving the rest of the object free.""" + duration: SpanDurationFilter + """Bound how long the span took, inclusive, in seconds.""" + + +_DURATION_FILTER_FIELDS = ("min", "max") + + +def _merge_deprecated_span_type(span_type: list[str] | None, filters: Any) -> Any: + """Fold the deprecated positional span_type argument into a filters object. + + Passing it both ways is an error rather than a merge, since there is no sensible way to + combine two lists that were each meant to be the whole constraint. span_type=[] keeps + its historical meaning of "no filter", which is why it cannot simply be copied across: + inside filters an empty list is rejected. + """ + if span_type is None: + return filters + warnings.warn( + "The span_type argument is deprecated; use filters={'span_type': [...]} instead.", + DeprecationWarning, + stacklevel=3, + ) + if filters is not None and not isinstance(filters, Mapping): + raise ValueError("filters must be an object") + if filters is not None and "span_type" in filters: + raise ValueError("span_type cannot be provided both directly and in filters") + if not span_type: + # Preserve the legacy behavior where span_type=[] means no filter. + return filters + return {**(filters or {}), "span_type": span_type} + + +def _btql_cmp(op: str, name: list[str], value: Any) -> dict[str, Any]: + """BTQL comparison between the column at path `name` and a literal value.""" + return {"op": op, "left": {"op": "ident", "name": name}, "right": {"op": "literal", "value": value}} + + +def _btql_null_check(op: str, name: list[str]) -> dict[str, Any]: + """BTQL "isnull" or "isnotnull" test on the column at path `name`.""" + return {"op": op, "expr": {"op": "ident", "name": name}} + + +class _FilterSpec(ABC): + """One field of SpanFilters, and everything the SDK knows about it. + + Each filter has to be understood twice: the server evaluates it as BTQL, and the SDK + evaluates it directly against spans it already holds. Those two readings have to agree, + so they sit on one class instead of in per-field branches spread across three distant + functions where they can quietly drift apart. + + Adding a filter means a new subclass plus a SpanFilters entry; + test_every_span_filter_field_has_a_spec is what keeps that pairing honest. + """ + + field: ClassVar[str] + """The SpanFilters key this class implements.""" + cacheable: ClassVar[bool] = False + """Whether CachedSpanFetcher's cache can answer this filter without asking the server. + + That cache is partitioned by span type and knows only that it holds every span of a + given type. For any other field it cannot tell a genuine empty result from a gap in + what it has fetched, so the query has to go out. + """ + + @abstractmethod + def validate(self, value: Any) -> Any: + """Check a raw filter value and return it in the form the other two methods expect. + + Raises ValueError if the value constrains nothing. An empty list or object is read + as a mistake rather than as a request for everything, so that a filter built up + programmatically fails loudly instead of quietly matching the whole trace. + """ + + @abstractmethod + def matches(self, span: Any, value: Any) -> bool: + """Evaluate this filter locally, against a SpanData or CachedSpan. + + `value` has already been through validate(). The result must agree with to_btql(): + a span the server would have returned is one this returns True for. + """ + + @abstractmethod + def to_btql(self, value: Any) -> list[dict[str, Any]]: + """Compile this filter into BTQL clauses, ANDed with the rest of the query. + + `value` has already been through validate(). Returning several clauses is normal -- + a tag filter emits one per tag -- but returning none would mean the filter + constrains nothing, which validate() is responsible for having rejected. + """ + + +class _SpanAttributeFilter(_FilterSpec): + """Base for filters that match one span_attributes key against a list of values. + + Matching is exact and case-sensitive, and the list is a set of alternatives: a span + matches if the attribute equals any entry. + """ + + attribute: ClassVar[str] + """The key to read out of span_attributes.""" + missing: ClassVar[Any] = None + """Stands in for the attribute when a span does not carry it at all.""" + + def validate(self, value: Any) -> list[str]: + if not isinstance(value, list) or not value or not all(isinstance(item, str) for item in value): + raise ValueError(f"filters.{self.field} must be a non-empty list of strings") + return list(value) + + def matches(self, span: Any, value: Any) -> bool: + return (getattr(span, "span_attributes", None) or {}).get(self.attribute, self.missing) in value + + def to_btql(self, value: Any) -> list[dict[str, Any]]: + return [_btql_cmp("in", ["span_attributes", self.attribute], value)] + + +class _SpanTypeFilter(_SpanAttributeFilter): + """Filter on span_attributes.type, the one field the span cache is partitioned by.""" + + field = "span_type" + attribute = "type" + # CachedSpanFetcher files typeless spans under "", so an explicit [""] query finds them. + missing = "" + cacheable = True + + +class _NameFilter(_SpanAttributeFilter): + """Filter on span_attributes.name. A span with no name matches no name filter.""" + + field = "name" + attribute = "name" + + +class _HasErrorFilter(_FilterSpec): + """Filter on whether the span recorded an error, testing only for presence. + + Any non-null error counts, whatever shape it has. + """ + + field = "has_error" + + def validate(self, value: Any) -> bool: + if not isinstance(value, bool): + raise ValueError("filters.has_error must be a boolean") + return value + + def matches(self, span: Any, value: Any) -> bool: + return (getattr(span, "error", None) is not None) == value + + def to_btql(self, value: Any) -> list[dict[str, Any]]: + return [_btql_null_check("isnotnull" if value else "isnull", ["error"])] + + +class _MetadataFilter(_FilterSpec): + """Filter on metadata by deep partial match. + + Only the keys named are compared, at any depth, so this narrows without having to + describe the whole metadata object. A None leaf matches a key whose value is null. + """ + + field = "metadata" + + def validate(self, value: Any, path: str = "filters.metadata") -> dict[str, Any]: + """Copy the filter, checking it at every depth. + + Empty objects are rejected wherever they appear, since they constrain nothing. + `path` follows the recursion so the error names the offending key rather than the + filter as a whole. + """ + if not isinstance(value, Mapping): + raise ValueError(f"{path} must be an object") + if not value: + raise ValueError(f"{path} must not be empty") + + result = {} + for key, child in value.items(): + if not isinstance(key, str) or not key: + raise ValueError(f"{path} keys must be non-empty strings") + result[key] = self.validate(child, f"{path}.{key}") if isinstance(child, Mapping) else child + return result + + def matches(self, span: Any, value: Any) -> bool: + return self._contains(getattr(span, "metadata", None), value) + + def _contains(self, actual: Any, expected: Mapping[str, Any]) -> bool: + """Whether `actual` carries every leaf of `expected`. + + Keys `actual` has and `expected` does not are ignored, so {"a": 1} matches a span + whose metadata is {"a": 1, "b": 2}. + """ + if not isinstance(actual, Mapping): + return False + for key, expected_value in expected.items(): + if key not in actual: + return False + if isinstance(expected_value, Mapping): + if not self._contains(actual[key], expected_value): + return False + elif actual[key] != expected_value: + return False + return True + + def to_btql(self, value: Any, path: tuple[str, ...] = ("metadata",)) -> list[dict[str, Any]]: + """Flatten the filter into one comparison per leaf. + + BTQL cannot match an object partially, so {"a": {"b": 1}} has to compile to a single + comparison against the column metadata.a.b. `path` accumulates that column path as + the recursion descends. + """ + children: list[dict[str, Any]] = [] + for key, child in value.items(): + child_path = (*path, key) + if isinstance(child, Mapping): + children.extend(self.to_btql(child, child_path)) + elif child is None: + children.append(_btql_null_check("isnull", list(child_path))) + else: + children.append(_btql_cmp("eq", list(child_path), child)) + return children + + +class _DurationFilter(_FilterSpec): + """Filter on elapsed wall-clock time, metrics.end - metrics.start, in seconds. + + Both bounds are inclusive. A span that is still open, or whose start/end metrics are + missing or non-numeric, has no duration and so matches no duration filter. + """ + + field = "duration" + + def validate(self, value: Any) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError("filters.duration must be an object") + unknown = set(value) - set(_DURATION_FILTER_FIELDS) + if unknown: + raise ValueError(f"Unsupported filters.duration fields: {', '.join(sorted(unknown))}") + # Iterating the constant rather than `value` keeps the bounds, and the BTQL built + # from them, in a fixed order. + bounds = {bound: self._bound(value[bound], bound) for bound in _DURATION_FILTER_FIELDS if bound in value} + if not bounds: + raise ValueError(f"filters.duration must specify at least one of: {', '.join(_DURATION_FILTER_FIELDS)}") + if "min" in bounds and "max" in bounds and bounds["min"] > bounds["max"]: + raise ValueError("filters.duration.min must be less than or equal to filters.duration.max") + return bounds + + def _bound(self, value: Any, name: str) -> float: + """Check one end of the range. + + NaN and infinities are rejected because they make a bound that can never be + satisfied or never be violated, and negatives because elapsed time cannot be + negative. + """ + if not self._is_real_number(value) or not math.isfinite(value) or value < 0: + raise ValueError(f"filters.duration.{name} must be a finite, non-negative number") + return value + + def matches(self, span: Any, value: Any) -> bool: + metrics = getattr(span, "metrics", None) + if not isinstance(metrics, Mapping): + return False + start, end = metrics.get("start"), metrics.get("end") + if not self._is_real_number(start) or not self._is_real_number(end): + return False + elapsed = end - start + if "min" in value and elapsed < value["min"]: + return False + if "max" in value and elapsed > value["max"]: + return False + return True + + @staticmethod + def _is_real_number(value: Any) -> bool: + """Whether `value` is a number that orders against a duration bound. + + is_numeric already rules out bool; complex is numeric but has no ordering. + """ + return is_numeric(value) and not isinstance(value, complex) + + def to_btql(self, value: Any) -> list[dict[str, Any]]: + elapsed = { + "op": "sub", + "left": {"op": "ident", "name": ["metrics", "end"]}, + "right": {"op": "ident", "name": ["metrics", "start"]}, + } + children = [] + if "min" in value: + children.append({"op": "ge", "left": elapsed, "right": {"op": "literal", "value": value["min"]}}) + if "max" in value: + children.append({"op": "le", "left": elapsed, "right": {"op": "literal", "value": value["max"]}}) + return children + + +# Registration order fixes the order of BTQL `and` children; keep it aligned with SpanFilters. +_FILTER_SPECS: dict[str, _FilterSpec] = { + spec.field: spec + for spec in ( + _SpanTypeFilter(), + _NameFilter(), + _HasErrorFilter(), + _MetadataFilter(), + _DurationFilter(), + ) +} +_FILTER_FIELDS = frozenset(_FILTER_SPECS) + + +def _normalize_span_filters(filters: Any) -> SpanFilters: + """Validate user-supplied filters into the form the rest of this module assumes. + + Every consumer downstream -- the local matcher, SpanFetcher's BTQL, CachedSpanFetcher's + reasoning about what its cache can answer -- takes filters as already normalized, so + this is the one place bad input has to be caught. Unknown or vacuous fields raise + rather than being dropped, since a silently ignored filter returns too many spans and + looks like a bug elsewhere. + """ + if filters is None: + return {} + if not isinstance(filters, Mapping): + raise ValueError("filters must be an object") + + unknown_fields = set(filters) - _FILTER_FIELDS + if unknown_fields: + raise ValueError(f"Unsupported span filter fields: {', '.join(sorted(unknown_fields))}") + + return cast( + SpanFilters, + {field: spec.validate(filters[field]) for field, spec in _FILTER_SPECS.items() if field in filters}, + ) + + +def _matches_span_filters(span: Any, filters: SpanFilters) -> bool: + """Evaluate normalized `filters` against a span the SDK already holds. + + This is the path taken for spans served out of a cache instead of fetched, so it has to + agree with the BTQL SpanFetcher would have sent for the same filters. + """ + if not filters: + return True + return all(_FILTER_SPECS[field].matches(span, value) for field, value in filters.items()) class SpanData: - """Span data returned by get_spans().""" + """One span, as returned by get_spans(). + + Fields mirror the span columns; anything the server sends that is not named explicitly + is still kept, as an attribute, so a newer backend does not lose data on the way through. + """ def __init__( self, @@ -49,16 +422,12 @@ def __init__( @classmethod def from_dict(cls, data: dict[str, Any]) -> "SpanData": - """Create SpanData from a dictionary.""" + """Build a span from a row, keeping columns this class does not name.""" return cls(**data) def to_dict(self) -> dict[str, Any]: - """Convert to dictionary.""" - result = {} - for key, value in self.__dict__.items(): - if value is not None: - result[key] = value - return result + """Return the span's set fields, dropping those left as None.""" + return clean_nones(self.__dict__) class SpanFetcher(ObjectFetcher[dict[str, Any]]): @@ -73,12 +442,12 @@ def __init__( object_id: str, root_span_id: str, state: BraintrustState, - span_type_filter: list[str] | None = None, include_scorers: bool = False, brainstore_realtime: bool = True, + filters: SpanFilters | None = None, ): - # 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) + # `filters` is expected to already be normalized by _normalize_span_filters. + filter_expr = self._build_filter(root_span_id, filters, include_scorers) super().__init__( object_type=object_type, @@ -91,52 +460,33 @@ def __init__( @staticmethod def _build_filter( root_span_id: str, - span_type_filter: list[str] | None = None, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> dict[str, Any]: - """Build BTQL filter expression.""" - children = [ - # Base filter: root_span_id = 'value' - { - "op": "eq", - "left": {"op": "ident", "name": ["root_span_id"]}, - "right": {"op": "literal", "value": root_span_id}, - }, - ] + """Compile the fetch into a single BTQL filter expression. + + Children are ANDed. Their order follows _FILTER_SPECS registration order and is + pinned by tests, so a new filter cannot silently reshuffle the query. + """ + # Scorer exclusion is a fetch mode rather than a SpanFilters field, so it stays here. + purpose = ["span_attributes", "purpose"] + children: list[dict[str, Any]] = [_btql_cmp("eq", ["root_span_id"], root_span_id)] if not include_scorers: children.append( { "op": "or", "children": [ - { - "op": "isnull", - "expr": { - "op": "ident", - "name": ["span_attributes", "purpose"], - }, - }, - { - "op": "ne", - "left": { - "op": "ident", - "name": ["span_attributes", "purpose"], - }, - "right": {"op": "literal", "value": "scorer"}, - }, + _btql_null_check("isnull", purpose), + _btql_cmp("ne", purpose, "scorer"), ], } ) - # If span type filter specified, add it - if span_type_filter and len(span_type_filter) > 0: - children.append( - { - "op": "in", - "left": {"op": "ident", "name": ["span_attributes", "type"]}, - "right": {"op": "literal", "value": span_type_filter}, - } - ) + filter_values: Mapping[str, Any] = filters or {} + for field, spec in _FILTER_SPECS.items(): + if field in filter_values: + children.extend(spec.to_btql(filter_values[field])) return {"op": "and", "children": children} @@ -148,8 +498,8 @@ def _get_state(self) -> BraintrustState: return self._state -SpanFetchFn = Callable[[list[str] | None], Awaitable[list[SpanData]]] -SpanFetchWithOptionsFn = Callable[[list[str] | None, bool], Awaitable[list[SpanData]]] +SpanFetchFn = Callable[[SpanFilters], Awaitable[list[SpanData]]] +SpanFetchWithOptionsFn = Callable[[SpanFilters, bool], Awaitable[list[SpanData]]] class GetThreadOptions(TypedDict, total=False): @@ -158,12 +508,14 @@ class GetThreadOptions(TypedDict, total=False): class CachedSpanFetcher: """ - Cached span fetcher that handles fetching and caching spans by type. - - Caching strategy: - - Cache spans by span type (dict[spanType, list[SpanData]]) - - Track if all spans have been fetched (all_fetched flag) - - When filtering by spanType, only fetch types not already in cache + Fetches spans for one root span, reusing what it has already seen. + + The cache is keyed by span type, plus a flag for whether an unfiltered fetch has + happened. That shape is what makes it useful and also what bounds it: it can answer a + span_type query offline, because it knows it holds every span of the types it has + fetched, but it cannot answer a query on any other field, because a partial result set + says nothing about the spans it never asked for. Those queries go to the server every + time and their results are used once rather than cached. """ def __init__( @@ -179,13 +531,14 @@ def __init__( self._all_fetched = False if fetch_fn is not None: - # Direct fetch function injection (for testing) + # Direct fetch function injection (for testing). Like the server, the injected + # function is responsible for honoring every filter it is given. async def _fetch_fn( - span_type: list[str] | None, + filters: SpanFilters, include_scorers: bool = False, ) -> list[SpanData]: del include_scorers - return await fetch_fn(span_type) + return await fetch_fn(filters) self._fetch_fn: SpanFetchWithOptionsFn = _fetch_fn else: @@ -196,7 +549,7 @@ async def _fetch_fn( ) async def _fetch_fn( - span_type: list[str] | None, + filters: SpanFilters, include_scorers: bool = False, ) -> list[SpanData]: state = await get_state() @@ -205,9 +558,9 @@ async def _fetch_fn( object_id=object_id, root_span_id=root_span_id, state=state, - span_type_filter=span_type, include_scorers=include_scorers, brainstore_realtime=brainstore_realtime, + filters=filters, ) rows = list(fetcher.fetch()) return [ @@ -237,29 +590,39 @@ async def _fetch_fn( async def get_spans( self, - span_type: list[str] | None = None, *, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> list[SpanData]: """ - Get spans, using cache when possible. + Get spans, using the cache where it can answer the query. Args: - span_type: Optional list of span types to filter by + filters: Optional filters for span type, name, error state, metadata, and duration include_scorers: Include spans with span_attributes.purpose = "scorer" Returns: List of matching spans """ + filters = _normalize_span_filters(filters) + span_type = filters.get("span_type") + # A partial cache is only authoritative for the fields it partitions on. + has_advanced_filters = any(not _FILTER_SPECS[field].cacheable for field in filters) + if include_scorers: - return await self._fetch_fn(span_type, True) + return await self._fetch_fn(filters, True) - # If we've fetched all spans, just filter from cache + # A complete cache can answer every supported filter locally. if self._all_fetched: - return self._get_from_cache(span_type) + spans = self._get_from_cache(span_type) + return [span for span in spans if _matches_span_filters(span, filters)] if has_advanced_filters else spans - # If no filter requested, fetch everything - if not span_type or len(span_type) == 0: + # Arbitrary filtered results are not authoritative for their span type. + if has_advanced_filters: + return await self._fetch_fn(filters, False) + + # If no filter requested, fetch everything. + if not span_type: # A full fetch is authoritative; reset the per-type cache first so a # prior typed fetch's spans are not duplicated by re-fetching them # (_fetch_spans appends). @@ -269,20 +632,19 @@ async def get_spans( self._all_fetched = True return self._get_from_cache(None) - # Find which spanTypes we don't have in cache yet + # Find which span types we don't have in cache yet. missing_types = [t for t in span_type if t not in self._span_cache] - - # If all requested types are cached, return from cache - if not missing_types: - return self._get_from_cache(span_type) - - # Fetch only the missing types - await self._fetch_spans(missing_types) + if missing_types: + await self._fetch_spans(missing_types) return self._get_from_cache(span_type) async def _fetch_spans(self, span_type: list[str] | None) -> None: - """Fetch spans from the server.""" - spans = await self._fetch_fn(span_type, False) + """Fetch spans and file them into the cache under their own type. + + Spans are filed by the type they report, not the type that was asked for, so a + requested type that yields nothing leaves no entry and will be asked for again. + """ + spans = await self._fetch_fn({"span_type": span_type} if span_type else {}, False) for span in spans: span_attrs = span.span_attributes or {} @@ -292,7 +654,11 @@ async def _fetch_spans(self, span_type: list[str] | None) -> None: self._span_cache[span_type_str].append(span) def _get_from_cache(self, span_type: list[str] | None) -> list[SpanData]: - """Get spans from cache, optionally filtering by type.""" + """Read spans back out of the cache, optionally narrowing to some types. + + Assumes the caller has established that the cache holds what is being asked for; + types with no entry are simply absent from the result, not fetched. + """ if not span_type or len(span_type) == 0: # Return all spans result = [] @@ -322,13 +688,15 @@ async def get_spans( self, span_type: list[str] | None = None, *, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> list[SpanData]: """ Fetch all spans for this root span. Args: - span_type: Optional list of span types to filter by + span_type: Deprecated; use filters["span_type"] instead + filters: Optional filters for span type, name, error state, metadata, and duration include_scorers: Include spans with span_attributes.purpose = "scorer" Returns: @@ -349,7 +717,7 @@ async def get_thread(self, options: GetThreadOptions | None = None) -> list[Any] ... -class LocalTrace(dict): +class LocalTrace(dict[str, Any]): """ SDK implementation of Trace that uses local span cache and falls back to BTQL. Carries identifying information about the evaluation so scorers can perform @@ -413,6 +781,7 @@ async def get_spans( self, span_type: list[str] | None = None, *, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> list[SpanData]: """ @@ -421,46 +790,29 @@ async def get_spans( back to CachedSpanFetcher which handles BTQL fetching and caching. Args: - span_type: Optional list of span types to filter by + span_type: Deprecated; use filters["span_type"] instead + filters: Optional filters for span type, name, error state, metadata, and duration include_scorers: Include spans with span_attributes.purpose = "scorer" Returns: List of matching spans """ + normalized_filters = _normalize_span_filters(_merge_deprecated_span_type(span_type, filters)) + # Try local span cache first (for recently logged spans not yet flushed) cached_spans = self._state.span_cache.get_by_root_span_id(self._root_span_id) if cached_spans and len(cached_spans) > 0: - # Filter by purpose spans = [ span for span in cached_spans - if include_scorers or not (span.span_attributes or {}).get("purpose") == "scorer" + if (include_scorers or not (span.span_attributes or {}).get("purpose") == "scorer") + and _matches_span_filters(span, normalized_filters) ] - # Filter by span type if requested - if span_type and len(span_type) > 0: - spans = [span for span in spans if (span.span_attributes or {}).get("type", "") in span_type] - - # Convert to SpanData - return [ - SpanData( - input=span.input, - output=span.output, - expected=getattr(span, "expected", None), - error=getattr(span, "error", None), - scores=getattr(span, "scores", None), - metrics=getattr(span, "metrics", None), - metadata=span.metadata, - span_id=span.span_id, - span_parents=span.span_parents, - span_attributes=span.span_attributes, - tags=getattr(span, "tags", None), - ) - for span in spans - ] + return [SpanData.from_dict(span.to_dict()) for span in spans] - # Fall back to CachedSpanFetcher for BTQL fetching with caching - return await self._cached_fetcher.get_spans(span_type, include_scorers=include_scorers) + # Fall back to CachedSpanFetcher for BTQL fetching with caching. + return await self._cached_fetcher.get_spans(filters=normalized_filters, include_scorers=include_scorers) async def get_thread(self, options: GetThreadOptions | None = None) -> list[Any]: """ @@ -479,7 +831,7 @@ async def _fetch_thread(self, options: GetThreadOptions | None = None) -> list[A await asyncio.get_event_loop().run_in_executor(None, lambda: self._state.login()) preprocessor = options.get("preprocessor") if options and options.get("preprocessor") else None - result = await asyncio.get_event_loop().run_in_executor( + result: Any = await asyncio.get_event_loop().run_in_executor( None, lambda: invoke( global_function=preprocessor or "project_default", @@ -498,15 +850,20 @@ async def _fetch_thread(self, options: GetThreadOptions | None = None) -> list[A return result if isinstance(result, list) else [] async def _ensure_spans_ready(self) -> None: - """Ensure spans are flushed before fetching.""" - if self._spans_flushed or not self._ensure_spans_flushed: + """Flush pending spans so a fetch sees them, at most once per trace. + + Concurrent scorers share one in-flight flush rather than each triggering their own. + A failed flush clears that shared handle so the next caller can retry. + """ + ensure_spans_flushed = self._ensure_spans_flushed + if self._spans_flushed or ensure_spans_flushed is None: return if self._spans_flush_promise is None: - async def flush_and_mark(): + async def flush_and_mark() -> None: try: - await self._ensure_spans_flushed() + await ensure_spans_flushed() self._spans_flushed = True except Exception as err: self._spans_flush_promise = None