From 10b2a3ea7d7e537ed7cc6169ff35f70911d5a7b0 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, "tags": {"all": ["production"], "none": ["internal"]}, "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 | 53 ++- py/src/braintrust/test_span_cache.py | 7 + py/src/braintrust/test_trace.py | 507 +++++++++++++------- py/src/braintrust/trace.py | 680 ++++++++++++++++++++++----- 5 files changed, 933 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..f30ec7190 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,10 @@ 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, + **kwargs: Any, ): self.span_id = span_id self.input = input @@ -40,33 +49,29 @@ def __init__( self.metadata = metadata self.span_parents = span_parents self.span_attributes = span_attributes + self.error = error + self.metrics = metrics + self.tags = tags + # Retain any fields written by a newer SDK so to_dict() round-trips them. + for key, value in kwargs.items(): + setattr(self, key, value) 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 its cached record, keeping fields this class does not name. + + Retaining unknown keys means a cache written by a newer SDK stays readable, and + round-trips back out through to_dict() intact. + """ + 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..ba77f7037 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")] - - 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 + 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. - @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" - - 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 + result = await fetcher.get_spans(filters={"span_type": span_type} if span_type else None) - 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,198 @@ 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, + "tags": { + "all": ["production"], + "any": ["priority", "customer-facing"], + "none": ["internal"], + }, + "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}} + + def includes(tag): + return { + "op": "includes", + "haystack": {"op": "ident", "name": ["tags"]}, + "needle": {"op": "literal", "value": tag}, + } + + 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"]), + includes("production"), + {"op": "or", "children": [includes("priority"), includes("customer-facing")]}, + {"op": "or", "children": [isnull(["tags"]), {"op": "not", "expr": includes("internal")}]}, + 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"}, + tags=["production", "priority"], + 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"}, + tags=["production"], + 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, + "tags": {"all": ["production"], "any": ["priority", "customer-facing"]}, + "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"), + ({"tags": {"all": []}}, "tags.all"), + ({"duration": {"min": -1}}, "duration.min"), + ({"duration": {"min": 2, "max": 1}}, "duration.min"), + # A field that is present must constrain something, at any nesting depth. + ({"tags": {}}, "filters.tags must specify at least one of: all, any, none"), + ({"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 +443,109 @@ 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}, + tags=["production", "priority"], + ) + + 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, + "tags": {"all": ["production"], "any": ["priority"]}, + "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 because no field was written, unlike filters={"tags": {}}.""" + 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..6e4ec0a22 100644 --- a/py/src/braintrust/trace.py +++ b/py/src/braintrust/trace.py @@ -6,16 +6,488 @@ """ 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 SpanTagFilter(TypedDict, total=False): + """Exact, case-sensitive tag matching options.""" + + all: list[str] + """Every tag must be present.""" + any: list[str] + """At least one tag must be present.""" + none: list[str] + """No tag may be present.""" + + +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.""" + tags: SpanTagFilter + """Require, allow, or exclude tags. The options given are ANDed together.""" + 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.""" + + +# Ordered so error messages and normalized dicts list sub-fields deterministically. +_TAG_FILTER_FIELDS = ("all", "any", "none") +_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 _validate_string_list(value: Any, field: str) -> list[str]: + """Copy `value` as a list of strings, rejecting empty lists and non-string entries.""" + if not isinstance(value, list) or not value or not all(isinstance(item, str) for item in value): + raise ValueError(f"filters.{field} must be a non-empty list of strings") + return list(value) + + +def _validate_metadata_filter(value: Any, path: str = "filters.metadata") -> dict[str, Any]: + """Copy a deep partial metadata filter, checking it at every depth. + + Empty objects are rejected wherever they appear, since they constrain nothing. `path` + tracks the position being checked 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] = _validate_metadata_filter(child, f"{path}.{key}") if isinstance(child, Mapping) else child + return result + + +def _as_number(value: Any) -> float | None: + """Return `value` if it is a real number, else None. + + bool is excluded despite subclassing int, and complex despite being numeric: neither + orders meaningfully against a duration bound. + """ + return value if is_numeric(value) and not isinstance(value, complex) else None + + +def _validate_duration_bound(value: Any, field: str) -> float: + """Check one end of a duration range, in seconds. + + 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. + """ + number = _as_number(value) + if number is None or not math.isfinite(number) or number < 0: + raise ValueError(f"filters.{field} must be a finite, non-negative number") + return number + + +def _validate_sub_filter( + value: Any, + field: str, + subfields: tuple[str, ...], + validate: Callable[[Any, str], Any], +) -> dict[str, Any]: + """Validate a nested filter object such as `tags` or `duration`. + + Requires at least one recognized sub-field: an object that is present but empty would + otherwise silently match every span. Results come back in `subfields` order, which is + what keeps the generated BTQL deterministic. + """ + if not isinstance(value, Mapping): + raise ValueError(f"filters.{field} must be an object") + unknown = set(value) - set(subfields) + if unknown: + raise ValueError(f"Unsupported filters.{field} fields: {', '.join(sorted(unknown))}") + normalized = {sub: validate(value[sub], f"{field}.{sub}") for sub in subfields if sub in value} + if not normalized: + raise ValueError(f"filters.{field} must specify at least one of: {', '.join(subfields)}") + return normalized + + +def _metadata_matches(actual: Any, expected: Mapping[str, Any]) -> bool: + """Whether `actual` contains every leaf of the deep partial filter `expected`. + + Keys `actual` has and `expected` does not are ignored, so {"a": 1} matches a span whose + metadata is {"a": 1, "b": 2}. Nested mappings recurse; other leaves compare with ==. + """ + if not isinstance(actual, Mapping): + return False + for key, expected_value in expected.items(): + if key not in actual: + return False + actual_value = actual[key] + if isinstance(expected_value, Mapping): + if not _metadata_matches(actual_value, expected_value): + return False + elif actual_value != expected_value: + return False + return True + + +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}} + + +def _btql_has_tag(tag: str) -> dict[str, Any]: + """BTQL test that the span's tags array contains `tag`.""" + return { + "op": "includes", + "haystack": {"op": "ident", "name": ["tags"]}, + "needle": {"op": "literal", "value": tag}, + } + + +def _btql_metadata_filters(metadata: Mapping[str, Any], path: tuple[str, ...] = ("metadata",)) -> list[dict[str, Any]]: + """Flatten a deep partial metadata 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. + """ + filters: list[dict[str, Any]] = [] + for key, value in metadata.items(): + child = (*path, key) + if isinstance(value, Mapping): + filters.extend(_btql_metadata_filters(value, child)) + elif value is None: + filters.append(_btql_null_check("isnull", list(child))) + else: + filters.append(_btql_cmp("eq", list(child), value)) + return filters + + +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]: + return _validate_string_list(value, self.field) + + 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 _TagsFilter(_FilterSpec): + """Filter on the span's tags by exact, case-sensitive membership. + + `all`, `any` and `none` may be combined and are ANDed together. An untagged span + satisfies `none` and fails both `all` and `any`. + """ + + field = "tags" + + def validate(self, value: Any) -> dict[str, Any]: + return _validate_sub_filter(value, self.field, _TAG_FILTER_FIELDS, _validate_string_list) + + def matches(self, span: Any, value: Any) -> bool: + tags = getattr(span, "tags", None) or [] + present = set(tags) if isinstance(tags, (list, tuple, set)) else set() + if "all" in value and not present.issuperset(value["all"]): + return False + if "any" in value and present.isdisjoint(value["any"]): + return False + if "none" in value and not present.isdisjoint(value["none"]): + return False + return True + + def to_btql(self, value: Any) -> list[dict[str, Any]]: + children = [_btql_has_tag(tag) for tag in value.get("all", [])] + any_tags = value.get("any", []) + if any_tags: + children.append({"op": "or", "children": [_btql_has_tag(tag) for tag in any_tags]}) + # An absent tags column is not "tagged with X", so it satisfies a `none` clause. + children.extend( + { + "op": "or", + "children": [ + _btql_null_check("isnull", ["tags"]), + {"op": "not", "expr": _btql_has_tag(tag)}, + ], + } + for tag in value.get("none", []) + ) + return children + + +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) -> dict[str, Any]: + return _validate_metadata_filter(value) + + def matches(self, span: Any, value: Any) -> bool: + return _metadata_matches(getattr(span, "metadata", None), value) + + def to_btql(self, value: Any) -> list[dict[str, Any]]: + return _btql_metadata_filters(value) + + +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]: + bounds = _validate_sub_filter(value, self.field, _DURATION_FILTER_FIELDS, _validate_duration_bound) + 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 matches(self, span: Any, value: Any) -> bool: + metrics = getattr(span, "metrics", None) + if not isinstance(metrics, Mapping): + return False + start = _as_number(metrics.get("start")) + end = _as_number(metrics.get("end")) + if start is None or end is None: + return False + duration = end - start + if "min" in value and duration < value["min"]: + return False + if "max" in value and duration > value["max"]: + return False + return True + + 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(), + _TagsFilter(), + _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 +521,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 +541,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 +559,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 +597,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 +607,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 +630,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 +648,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 +657,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 +689,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, tags, 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 +731,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 +753,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 +787,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, tags, metadata, and duration include_scorers: Include spans with span_attributes.purpose = "scorer" Returns: @@ -349,7 +816,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 +880,7 @@ async def get_spans( self, span_type: list[str] | None = None, *, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> list[SpanData]: """ @@ -421,46 +889,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, tags, 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 +930,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 +949,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