diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index ee4ba677..d1752d4b 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 ee926614..e3541419 100644 --- a/py/src/braintrust/span_cache.py +++ b/py/src/braintrust/span_cache.py @@ -33,6 +33,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,10 +43,13 @@ 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} + result: dict[str, Any] = {"span_id": self.span_id} if self.input is not None: result["input"] = self.input if self.output is not None: @@ -54,6 +60,12 @@ def to_dict(self) -> dict[str, Any]: result["span_parents"] = self.span_parents if self.span_attributes is not None: result["span_attributes"] = self.span_attributes + if self.error is not None: + result["error"] = self.error + if self.metrics is not None: + result["metrics"] = self.metrics + if self.tags is not None: + result["tags"] = self.tags return result @classmethod @@ -66,6 +78,9 @@ def from_dict(cls, data: dict[str, Any]) -> "CachedSpan": metadata=data.get("metadata"), span_parents=data.get("span_parents"), span_attributes=data.get("span_attributes"), + error=data.get("error"), + metrics=data.get("metrics"), + tags=data.get("tags"), ) diff --git a/py/src/braintrust/test_span_cache.py b/py/src/braintrust/test_span_cache.py index 9b250d44..767d169a 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 3a572f68..43ea8308 100644 --- a/py/src/braintrust/test_trace.py +++ b/py/src/braintrust/test_trace.py @@ -1,16 +1,20 @@ """Tests for Trace functionality.""" +import asyncio +from types import SimpleNamespace + import pytest from braintrust.trace import CachedSpanFetcher, LocalTrace, SpanData, SpanFetcher # Helper to create mock spans def make_span(span_id: str, span_type: str, **extra) -> SpanData: + span_attributes = extra.pop("span_attributes", {"type": span_type}) 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, ) @@ -56,7 +60,7 @@ async def fetch_fn(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] @@ -106,7 +110,7 @@ async def fetch_fn(span_type): return llm_spans fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - result = await fetcher.get_spans(span_type=["llm"]) + result = await fetcher.get_spans(filters={"span_type": ["llm"]}) assert call_count == 1 assert len(result) == 2 @@ -152,11 +156,11 @@ async def fetch_fn(span_type): fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) # First call - fetches llm spans - await fetcher.get_spans(span_type=["llm"]) + await fetcher.get_spans(filters={"span_type": ["llm"]}) assert call_count == 1 # Second call for same type - should use cache - result = await fetcher.get_spans(span_type=["llm"]) + result = await fetcher.get_spans(filters={"span_type": ["llm"]}) assert call_count == 1 # Still 1 assert len(result) == 2 @@ -180,11 +184,11 @@ async def fetch_fn(span_type): fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) # First call - fetches llm spans - await fetcher.get_spans(span_type=["llm"]) + await fetcher.get_spans(filters={"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"]) + result = await fetcher.get_spans(filters={"span_type": ["llm", "function"]}) assert call_count == 2 assert len(result) == 2 @@ -211,12 +215,12 @@ async def fetch_fn(span_type): assert call_count == 1 # Subsequent filtered calls should use cache - llm_result = await fetcher.get_spans(span_type=["llm"]) + llm_result = await fetcher.get_spans(filters={"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"]) + function_result = await fetcher.get_spans(filters={"span_type": ["function"]}) assert call_count == 1 # Still 1 assert len(function_result) == 1 assert function_result[0].span_id == "span-2" @@ -240,7 +244,7 @@ async def fetch_fn(span_type): await fetcher.get_spans() # Filter for llm and tool - result = await fetcher.get_spans(span_type=["llm", "tool"]) + result = await fetcher.get_spans(filters={"span_type": ["llm", "tool"]}) assert len(result) == 3 assert {s.span_id for s in result} == {"span-1", "span-3", "span-4"} @@ -258,7 +262,7 @@ async def fetch_fn(span_type): await fetcher.get_spans() # Query for non-existent type - result = await fetcher.get_spans(span_type=["nonexistent"]) + result = await fetcher.get_spans(filters={"span_type": ["nonexistent"]}) assert len(result) == 0 @pytest.mark.asyncio @@ -280,7 +284,7 @@ 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.asyncio @@ -324,11 +328,11 @@ async def fetch_fn(span_type): fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) # First call with type filter returns empty - result1 = await fetcher.get_spans(span_type=["llm"]) + result1 = await fetcher.get_spans(filters={"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"]) + result2 = await fetcher.get_spans(filters={"span_type": ["llm"]}) assert call_count == 2 assert len(result2) == 1 @@ -345,11 +349,208 @@ async def fetch_fn(span_type): fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) - result = await fetcher.get_spans(span_type=[]) + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + result = await fetcher.get_spans(span_type=[]) assert call_args[0] is None or call_args[0] == [] assert len(result) == 1 + 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", "request": {"region": "us-east-1"}}, + "duration": {"min": 0.5, "max": 10}, + }, + ) + + assert list(fetcher.fetch()) == [] + query_filter = calls[0]["json"]["query"]["filter"] + children = query_filter["children"] + + def comparison(op, path, value): + return { + "op": op, + "left": {"op": "ident", "name": path}, + "right": {"op": "literal", "value": value}, + } + + 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"]}, + } + assert query_filter["op"] == "and" + assert len(children) == 12 + assert comparison("eq", ["root_span_id"], "root-1") in children + assert { + "op": "or", + "children": [ + {"op": "isnull", "expr": {"op": "ident", "name": ["span_attributes", "purpose"]}}, + comparison("ne", ["span_attributes", "purpose"], "scorer"), + ], + } in children + assert comparison("in", ["span_attributes", "type"], ["tool"]) in children + assert comparison("in", ["span_attributes", "name"], ["search", "lookup"]) in children + assert {"op": "isnull", "expr": {"op": "ident", "name": ["error"]}} in children + assert includes("production") in children + assert {"op": "or", "children": [includes("priority"), includes("customer-facing")]} in children + assert { + "op": "or", + "children": [ + {"op": "isnull", "expr": {"op": "ident", "name": ["tags"]}}, + {"op": "not", "expr": includes("internal")}, + ], + } in children + assert comparison("eq", ["metadata", "model"], "gpt-5") in children + assert comparison("eq", ["metadata", "request", "region"], "us-east-1") in children + assert {"op": "ge", "left": duration, "right": {"op": "literal", "value": 0.5}} in children + assert {"op": "le", "left": duration, "right": {"op": "literal", "value": 10}} in children + + @pytest.mark.asyncio + async def test_advanced_filters_use_full_cache_when_available(self): + spans = [ + make_span( + "matching", + "tool", + error={"message": "boom"}, + tags=["production", "priority"], + metadata={"model": "gpt-5", "request": {"region": "us-east-1", "id": 1}}, + metrics={"start": 1, "end": 4}, + span_attributes={"type": "tool", "name": "search"}, + ), + make_span( + "too-fast", + "tool", + error={"message": "boom"}, + tags=["production"], + metadata={"model": "gpt-5"}, + metrics={"start": 1, "end": 1.1}, + span_attributes={"type": "tool", "name": "search"}, + ), + make_span("successful", "tool", span_attributes={"type": "tool", "name": "search"}), + ] + call_count = 0 + + async def fetch_fn(span_type): + nonlocal call_count + del span_type + call_count += 1 + 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}, + } + ) + + assert call_count == 1 + assert [span.span_id for span in result] == ["matching"] + + @pytest.mark.asyncio + async def test_advanced_filter_cache_coalesces_and_keys_empty_results(self): + started = asyncio.Event() + release = asyncio.Event() + call_count = 0 + spans = [make_span("search", "tool", span_attributes={"type": "tool", "name": "search"})] + + async def fetch_fn(span_type): + nonlocal call_count + assert span_type == ["tool"] + call_count += 1 + started.set() + await release.wait() + return spans + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + missing_filters = {"span_type": ["tool"], "name": ["missing"]} + first = asyncio.create_task(fetcher.get_spans(filters=missing_filters)) + await started.wait() + second = asyncio.create_task(fetcher.get_spans(filters=missing_filters)) + release.set() + + assert await first == [] + assert await second == [] + assert await fetcher.get_spans(filters=missing_filters) == [] + + search_filters = {"span_type": ["tool"], "name": ["search"]} + assert [span.span_id for span in await fetcher.get_spans(filters=search_filters)] == ["search"] + assert [span.span_id for span in await fetcher.get_spans(filters=search_filters)] == ["search"] + assert call_count == 2 + + @pytest.mark.asyncio + async def test_advanced_filter_cache_retries_failures(self): + call_count = 0 + + async def fetch_fn(span_type): + nonlocal call_count + assert span_type == ["tool"] + call_count += 1 + if call_count == 1: + raise RuntimeError("temporary failure") + return [] + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + filters = {"span_type": ["tool"], "name": ["missing"]} + + with pytest.raises(RuntimeError, match="temporary failure"): + await fetcher.get_spans(filters=filters) + + assert await fetcher.get_spans(filters=filters) == [] + assert await fetcher.get_spans(filters=filters) == [] + assert call_count == 2 + + @pytest.mark.asyncio + async def test_rejects_legacy_and_nested_span_type(self): + fetcher = CachedSpanFetcher(fetch_fn=lambda span_type: None) + + with pytest.warns(DeprecationWarning, match="span_type argument is deprecated"): + with pytest.raises(ValueError, match="span_type"): + await fetcher.get_spans(["llm"], filters={"span_type": ["tool"]}) + + @pytest.mark.parametrize( + ("filters", "message"), + [ + ({"name": []}, "name"), + ({"tags": {"all": []}}, "tags.all"), + ({"duration": {"min": -1}}, "duration.min"), + ({"duration": {"min": 2, "max": 1}}, "duration.min"), + ], + ) + @pytest.mark.asyncio + async def test_rejects_invalid_advanced_filters(self, filters, message): + fetcher = CachedSpanFetcher(fetch_fn=lambda span_type: None) + + with pytest.raises(ValueError, match=message): + await fetcher.get_spans(filters=filters) + @pytest.mark.parametrize( ("brainstore_realtime", "expected"), [ @@ -393,14 +594,69 @@ 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): + matching = SimpleNamespace( + input={"query": "weather"}, + output={"result": "sunny"}, + expected=None, + error={"message": "retryable"}, + scores=None, + metrics={"start": 1, "end": 4}, + metadata={"request": {"region": "us-east-1", "id": 1}}, + span_id="matching", + span_parents=[], + span_attributes={"type": "tool", "name": "search"}, + tags=["production", "priority"], + ) + wrong_name = SimpleNamespace( + **{ + **matching.__dict__, + "span_id": "wrong-name", + "span_attributes": {"type": "tool", "name": "lookup"}, + } + ) + scorer = SimpleNamespace( + **{ + **matching.__dict__, + "span_id": "scorer", + "span_attributes": {"type": "tool", "name": "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=[matching, wrong_name, scorer]), + ) + + 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"] + + 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 b72b8d62..c0ab0335 100644 --- a/py/src/braintrust/trace.py +++ b/py/src/braintrust/trace.py @@ -6,14 +6,245 @@ """ import asyncio -from collections.abc import Awaitable, Callable +import math +import warnings +from collections.abc import Awaitable, Callable, Coroutine, Mapping from typing import Any, Protocol, TypedDict +from braintrust.bt_json import bt_dumps from braintrust.functions.invoke import invoke from braintrust.logger import BraintrustState, ObjectFetcher from braintrust.types import Metadata +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.""" + + span_type: list[str] + """Exact span_attributes.type values to match with OR semantics.""" + name: list[str] + """Exact span_attributes.name values to match with OR semantics.""" + has_error: bool + """Whether the span's error field must be non-null.""" + tags: SpanTagFilter + """Exact tag matching options.""" + metadata: dict[str, Any] + """A deep partial metadata object whose leaves are matched by exact equality.""" + duration: SpanDurationFilter + """Inclusive duration bounds, in seconds.""" + + +_FILTER_FIELDS = {"span_type", "name", "has_error", "tags", "metadata", "duration"} +_TAG_FILTER_FIELDS = {"all", "any", "none"} +_DURATION_FILTER_FIELDS = {"min", "max"} + + +def _warn_deprecated_span_type(span_type: list[str] | None) -> None: + if span_type is not None: + warnings.warn( + "The span_type argument is deprecated; use filters={'span_type': [...]} instead.", + DeprecationWarning, + stacklevel=3, + ) + + +def _validate_string_list(value: Any, field: str) -> list[str]: + 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]: + if not isinstance(value, Mapping): + raise ValueError(f"{path} must be an object") + + 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") + if isinstance(child, Mapping): + normalized_child = _validate_metadata_filter(child, f"{path}.{key}") + if normalized_child: + result[key] = normalized_child + else: + result[key] = child + return result + + +def _validate_duration_bound(value: Any, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: + raise ValueError(f"filters.duration.{field} must be a finite, non-negative number") + return value + + +def _normalize_span_filters( + span_type: list[str] | None, + filters: Any, +) -> SpanFilters: + if filters is not None and not isinstance(filters, Mapping): + raise ValueError("filters must be an object") + + raw_filters = dict(filters or {}) + unknown_fields = set(raw_filters) - _FILTER_FIELDS + if unknown_fields: + raise ValueError(f"Unsupported span filter fields: {', '.join(sorted(unknown_fields))}") + if span_type is not None and "span_type" in raw_filters: + raise ValueError("span_type cannot be provided both directly and in filters") + + normalized: SpanFilters = {} + if span_type: + # Preserve the legacy behavior where span_type=[] means no filter. + normalized["span_type"] = _validate_string_list(span_type, "span_type") + elif "span_type" in raw_filters: + normalized["span_type"] = _validate_string_list(raw_filters["span_type"], "span_type") + + if "name" in raw_filters: + normalized["name"] = _validate_string_list(raw_filters["name"], "name") + + if "has_error" in raw_filters: + has_error = raw_filters["has_error"] + if not isinstance(has_error, bool): + raise ValueError("filters.has_error must be a boolean") + normalized["has_error"] = has_error + + if "tags" in raw_filters: + tags = raw_filters["tags"] + if not isinstance(tags, Mapping): + raise ValueError("filters.tags must be an object") + unknown_tag_fields = set(tags) - _TAG_FILTER_FIELDS + if unknown_tag_fields: + raise ValueError(f"Unsupported tag filter fields: {', '.join(sorted(unknown_tag_fields))}") + normalized_tags: SpanTagFilter = {} + for mode in ("all", "any", "none"): + if mode in tags: + normalized_tags[mode] = _validate_string_list(tags[mode], f"tags.{mode}") + if normalized_tags: + normalized["tags"] = normalized_tags + + if "metadata" in raw_filters: + normalized_metadata = _validate_metadata_filter(raw_filters["metadata"]) + if normalized_metadata: + normalized["metadata"] = normalized_metadata + + if "duration" in raw_filters: + duration = raw_filters["duration"] + if not isinstance(duration, Mapping): + raise ValueError("filters.duration must be an object") + unknown_duration_fields = set(duration) - _DURATION_FILTER_FIELDS + if unknown_duration_fields: + raise ValueError(f"Unsupported duration filter fields: {', '.join(sorted(unknown_duration_fields))}") + normalized_duration: SpanDurationFilter = {} + for bound in ("min", "max"): + if bound in duration: + normalized_duration[bound] = _validate_duration_bound(duration[bound], bound) + if ( + "min" in normalized_duration + and "max" in normalized_duration + and normalized_duration["min"] > normalized_duration["max"] + ): + raise ValueError("filters.duration.min must be less than or equal to filters.duration.max") + if normalized_duration: + normalized["duration"] = normalized_duration + + return normalized + + +def _get_span_field(span: Any, field: str) -> Any: + if isinstance(span, Mapping): + return span.get(field) + return getattr(span, field, None) + + +def _metadata_matches(actual: Any, expected: Mapping[str, Any]) -> bool: + 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 _matches_span_filters(span: Any, filters: SpanFilters) -> bool: + span_attributes = _get_span_field(span, "span_attributes") or {} + if not isinstance(span_attributes, Mapping): + span_attributes = {} + + span_types = filters.get("span_type") + if span_types is not None and span_attributes.get("type", "") not in span_types: + return False + + names = filters.get("name") + if names is not None and span_attributes.get("name") not in names: + return False + + has_error = filters.get("has_error") + if has_error is not None and (_get_span_field(span, "error") is not None) != has_error: + return False + + tag_filter = filters.get("tags") + if tag_filter is not None: + tags = _get_span_field(span, "tags") or [] + tag_set = set(tags) if isinstance(tags, (list, tuple, set)) else set() + if "all" in tag_filter and not all(tag in tag_set for tag in tag_filter["all"]): + return False + if "any" in tag_filter and not any(tag in tag_set for tag in tag_filter["any"]): + return False + if "none" in tag_filter and any(tag in tag_set for tag in tag_filter["none"]): + return False + + metadata_filter = filters.get("metadata") + if metadata_filter is not None and not _metadata_matches(_get_span_field(span, "metadata"), metadata_filter): + return False + + duration_filter = filters.get("duration") + if duration_filter: + metrics = _get_span_field(span, "metrics") + if not isinstance(metrics, Mapping): + return False + start = metrics.get("start") + end = metrics.get("end") + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, (int, float)) + or not isinstance(end, (int, float)) + ): + return False + duration = end - start + if "min" in duration_filter and duration < duration_filter["min"]: + return False + if "max" in duration_filter and duration > duration_filter["max"]: + return False + + return True + + class SpanData: """Span data returned by get_spans().""" @@ -76,9 +307,10 @@ def __init__( 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) + normalized_filters = _normalize_span_filters(span_type_filter, filters) + filter_expr = self._build_filter(root_span_id, normalized_filters, include_scorers) super().__init__( object_type=object_type, @@ -91,11 +323,11 @@ 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 = [ + children: list[dict[str, Any]] = [ # Base filter: root_span_id = 'value' { "op": "eq", @@ -128,18 +360,107 @@ def _build_filter( } ) - # If span type filter specified, add it - if span_type_filter and len(span_type_filter) > 0: + filters = filters or {} + if "span_type" in filters: children.append( { "op": "in", "left": {"op": "ident", "name": ["span_attributes", "type"]}, - "right": {"op": "literal", "value": span_type_filter}, + "right": {"op": "literal", "value": filters["span_type"]}, + } + ) + + if "name" in filters: + children.append( + { + "op": "in", + "left": {"op": "ident", "name": ["span_attributes", "name"]}, + "right": {"op": "literal", "value": filters["name"]}, + } + ) + + if "has_error" in filters: + children.append( + { + "op": "isnotnull" if filters["has_error"] else "isnull", + "expr": {"op": "ident", "name": ["error"]}, + } + ) + + tag_filter = filters.get("tags", {}) + for tag in tag_filter.get("all", []): + children.append(SpanFetcher._tag_filter(tag)) + any_tags = tag_filter.get("any", []) + if any_tags: + children.append( + { + "op": "or", + "children": [SpanFetcher._tag_filter(tag) for tag in any_tags], + } + ) + for tag in tag_filter.get("none", []): + children.append( + { + "op": "or", + "children": [ + {"op": "isnull", "expr": {"op": "ident", "name": ["tags"]}}, + {"op": "not", "expr": SpanFetcher._tag_filter(tag)}, + ], + } + ) + + children.extend(SpanFetcher._metadata_filters(filters.get("metadata", {}))) + + duration_filter = filters.get("duration", {}) + duration_expr = { + "op": "sub", + "left": {"op": "ident", "name": ["metrics", "end"]}, + "right": {"op": "ident", "name": ["metrics", "start"]}, + } + if "min" in duration_filter: + children.append( + { + "op": "ge", + "left": duration_expr, + "right": {"op": "literal", "value": duration_filter["min"]}, + } + ) + if "max" in duration_filter: + children.append( + { + "op": "le", + "left": duration_expr, + "right": {"op": "literal", "value": duration_filter["max"]}, } ) return {"op": "and", "children": children} + @staticmethod + def _tag_filter(tag: str) -> dict[str, Any]: + return { + "op": "includes", + "haystack": {"op": "ident", "name": ["tags"]}, + "needle": {"op": "literal", "value": tag}, + } + + @staticmethod + def _metadata_filters(metadata: Mapping[str, Any], path: tuple[str, ...] = ()) -> list[dict[str, Any]]: + filters: list[dict[str, Any]] = [] + for key, value in metadata.items(): + child_path = (*path, key) + if isinstance(value, Mapping): + filters.extend(SpanFetcher._metadata_filters(value, child_path)) + else: + filters.append( + { + "op": "eq", + "left": {"op": "ident", "name": ["metadata", *child_path]}, + "right": {"op": "literal", "value": value}, + } + ) + return filters + @property def id(self) -> str: return self._object_id @@ -149,7 +470,7 @@ def _get_state(self) -> BraintrustState: SpanFetchFn = Callable[[list[str] | None], Awaitable[list[SpanData]]] -SpanFetchWithOptionsFn = Callable[[list[str] | None, bool], Awaitable[list[SpanData]]] +SpanFetchWithOptionsFn = Callable[[SpanFilters, bool], Coroutine[Any, Any, list[SpanData]]] class GetThreadOptions(TypedDict, total=False): @@ -164,6 +485,7 @@ class CachedSpanFetcher: - 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 + - Cache and coalesce advanced requests by their canonical serialized filters """ def __init__( @@ -176,16 +498,18 @@ def __init__( brainstore_realtime: bool = True, ): self._span_cache: dict[str, list[SpanData]] = {} + self._filtered_cache: dict[str, asyncio.Task[list[SpanData]]] = {} self._all_fetched = False if fetch_fn is not None: # Direct fetch function injection (for testing) 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) + spans = await fetch_fn(filters.get("span_type")) + return [span for span in spans if _matches_span_filters(span, filters)] self._fetch_fn: SpanFetchWithOptionsFn = _fetch_fn else: @@ -196,7 +520,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 +529,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 [ @@ -239,27 +563,56 @@ 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. 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 """ + _warn_deprecated_span_type(span_type) + normalized_filters = _normalize_span_filters(span_type, filters) + return await self._get_spans(normalized_filters, include_scorers) + + async def _get_spans( + self, + filters: SpanFilters, + include_scorers: bool, + ) -> list[SpanData]: + normalized_span_type = filters.get("span_type") + advanced_filters = set(filters) - {"span_type"} + 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) - - # If no filter requested, fetch everything - if not span_type or len(span_type) == 0: + return [span for span in self._get_from_cache(None) if _matches_span_filters(span, filters)] + + # Arbitrary filtered results are not authoritative for their span type, + # so cache them only by their complete normalized filter. + if advanced_filters: + cache_key = bt_dumps(filters) + task = self._filtered_cache.get(cache_key) + if task is None: + task = asyncio.create_task(self._fetch_fn(filters, False)) + self._filtered_cache[cache_key] = task + try: + return list(await task) + except BaseException: + if self._filtered_cache.get(cache_key) is task: + del self._filtered_cache[cache_key] + raise + + # If no filter requested, fetch everything. + if not normalized_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 +622,22 @@ async def get_spans( self._all_fetched = True return self._get_from_cache(None) - # Find which spanTypes 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 + # Find which span types we don't have in cache yet. + missing_types = [ + span_type_value for span_type_value in normalized_span_type if span_type_value not in self._span_cache + ] if not missing_types: - return self._get_from_cache(span_type) + return self._get_from_cache(normalized_span_type) - # Fetch only the missing types await self._fetch_spans(missing_types) - return self._get_from_cache(span_type) + return self._get_from_cache(normalized_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) + filters: SpanFilters = {} + if span_type: + filters["span_type"] = span_type + spans = await self._fetch_fn(filters, False) for span in spans: span_attrs = span.span_attributes or {} @@ -322,13 +677,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 +706,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 +770,7 @@ async def get_spans( self, span_type: list[str] | None = None, *, + filters: SpanFilters | None = None, include_scorers: bool = False, ) -> list[SpanData]: """ @@ -421,27 +779,26 @@ 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 """ + _warn_deprecated_span_type(span_type) + normalized_filters = _normalize_span_filters(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, @@ -459,8 +816,8 @@ async def get_spans( 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(normalized_filters, include_scorers) async def get_thread(self, options: GetThreadOptions | None = None) -> list[Any]: """ @@ -479,7 +836,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", @@ -499,14 +856,15 @@ async def _fetch_thread(self, options: GetThreadOptions | None = None) -> list[A async def _ensure_spans_ready(self) -> None: """Ensure spans are flushed before fetching.""" - if self._spans_flushed or not self._ensure_spans_flushed: + 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 diff --git a/py/src/braintrust/type_tests/test_trace_filters.py b/py/src/braintrust/type_tests/test_trace_filters.py new file mode 100644 index 00000000..8e554c01 --- /dev/null +++ b/py/src/braintrust/type_tests/test_trace_filters.py @@ -0,0 +1,22 @@ +from braintrust.trace import SpanFilters, Trace + + +async def use_trace_filters(trace: Trace) -> None: + filters: SpanFilters = { + "span_type": ["llm", "tool"], + "name": ["search", "lookup"], + "has_error": False, + "tags": { + "all": ["production"], + "any": ["priority", "customer-facing"], + "none": ["internal"], + }, + "metadata": { + "model": "gpt-5", + "request": {"region": "us-east-1"}, + }, + "duration": {"min": 0.5, "max": 10}, + } + + await trace.get_spans(filters=filters) + await trace.get_spans(span_type=["llm"])