braintrust 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

braintrust/test_span_cache.py

@@ -0,0 +1,344 @@
+"""Tests for SpanCache (disk-based cache)."""
+
+
+from braintrust.span_cache import CachedSpan, SpanCache
+
+
+def test_span_cache_write_and_read():
+    """Test storing and retrieving spans by rootSpanId."""
+    cache = SpanCache()
+    cache.start()  # Start for testing (cache is disabled by default)
+
+    root_span_id = "root-123"
+    span1 = CachedSpan(
+        span_id="span-1",
+        input={"text": "hello"},
+        output={"response": "world"},
+    )
+    span2 = CachedSpan(
+        span_id="span-2",
+        input={"text": "foo"},
+        output={"response": "bar"},
+    )
+
+    cache.queue_write(root_span_id, span1.span_id, span1)
+    cache.queue_write(root_span_id, span2.span_id, span2)
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert len(spans) == 2
+
+    span_ids = {s.span_id for s in spans}
+    assert "span-1" in span_ids
+    assert "span-2" in span_ids
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_return_none_for_unknown():
+    """Test that unknown rootSpanId returns None."""
+    cache = SpanCache()
+    cache.start()
+
+    spans = cache.get_by_root_span_id("nonexistent")
+    assert spans is None
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_merge_on_duplicate_writes():
+    """Test that subsequent writes to same spanId merge data."""
+    cache = SpanCache()
+    cache.start()
+
+    root_span_id = "root-123"
+    span_id = "span-1"
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, input={"text": "hello"}),
+    )
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, output={"response": "world"}),
+    )
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert len(spans) == 1
+    assert spans[0].span_id == span_id
+    assert spans[0].input == {"text": "hello"}
+    assert spans[0].output == {"response": "world"}
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_merge_metadata():
+    """Test that metadata objects are merged."""
+    cache = SpanCache()
+    cache.start()
+
+    root_span_id = "root-123"
+    span_id = "span-1"
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, metadata={"key1": "value1"}),
+    )
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, metadata={"key2": "value2"}),
+    )
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert spans[0].metadata == {"key1": "value1", "key2": "value2"}
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_has():
+    """Test the has() method."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-123", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.has("root-123") is True
+    assert cache.has("nonexistent") is False
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_clear():
+    """Test clearing spans for a specific rootSpanId."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+
+    cache.clear("root-1")
+
+    assert cache.has("root-1") is False
+    assert cache.has("root-2") is True
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_clear_all():
+    """Test clearing all cached spans."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+
+    cache.clear_all()
+
+    assert cache.size == 0
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_size():
+    """Test the size property."""
+    cache = SpanCache()
+    cache.start()
+
+    assert cache.size == 0
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.queue_write("root-1", "span-2", CachedSpan(span_id="span-2"))  # Same root
+    assert cache.size == 1
+
+    cache.queue_write("root-2", "span-3", CachedSpan(span_id="span-3"))  # Different root
+    assert cache.size == 2
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_dispose():
+    """Test that dispose cleans up and allows reuse."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    # Stop first to decrement refcount, then dispose
+    cache.stop()
+    cache.dispose()
+
+    assert cache.size == 0
+    assert cache.has("root-1") is False
+
+    # Should be able to write again after dispose (if we start again)
+    cache.start()
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disable():
+    """Test that disable() prevents writes."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.disable()
+
+    # Writes after disable should be no-ops
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1  # Still 1, not 2
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disabled_getter():
+    """Test the disabled property."""
+    # Cache is disabled by default until start() is called
+    cache = SpanCache()
+    assert cache.disabled is True
+
+    cache.start()
+    assert cache.disabled is False
+
+    cache.disable()
+    assert cache.disabled is True
+
+    cache.dispose()
+
+
+def test_span_cache_disabled_from_constructor():
+    """Test that cache can be disabled via constructor."""
+    cache = SpanCache(disabled=True)
+    assert cache.disabled is True
+
+    # Writes should be no-ops
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 0
+    assert cache.get_by_root_span_id("root-1") is None
+
+    cache.dispose()
+
+
+def test_span_cache_start_stop_lifecycle():
+    """Test that stop() allows start() to work again."""
+    cache = SpanCache()
+
+    # Initially disabled by default
+    assert cache.disabled is True
+
+    # Start for first "eval"
+    cache.start()
+    assert cache.disabled is False
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    # Stop after first "eval"
+    cache.stop()
+    cache.dispose()
+    assert cache.disabled is True
+
+    # Start for second "eval" - should work!
+    cache.start()
+    assert cache.disabled is False
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disable_prevents_start():
+    """Test that disable() prevents start() from working."""
+    cache = SpanCache()
+
+    # Simulate disable being called
+    cache.disable()
+    assert cache.disabled is True
+
+    # start() should be a no-op after disable()
+    cache.start()
+    assert cache.disabled is True
+
+    # Writes should still be no-ops
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 0
+
+    cache.dispose()
+
+
+def test_span_cache_parallel_eval_refcount():
+    """Test reference counting for parallel evals."""
+    cache = SpanCache()
+
+    # Simulate two evals starting
+    cache.start()  # Eval 1
+    assert cache.disabled is False
+
+    cache.start()  # Eval 2
+    assert cache.disabled is False
+
+    # Write data from both evals
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 2
+
+    # Eval 1 finishes first
+    cache.dispose()  # Should NOT dispose (refcount = 2)
+    cache.stop()  # Decrements to 1
+
+    # Cache should still be enabled and data intact
+    assert cache.disabled is False
+    assert cache.size == 2
+    assert cache.get_by_root_span_id("root-1") is not None
+    assert cache.get_by_root_span_id("root-2") is not None
+
+    # Eval 2 finishes
+    cache.dispose()  # Should NOT dispose yet (refcount = 1)
+    cache.stop()  # Decrements to 0, disables cache
+
+    # Now cache should be disabled
+    assert cache.disabled is True
+
+    # Final dispose should now work
+    cache.dispose()  # NOW it disposes (refcount = 0)
+    assert cache.size == 0
+
+
+def test_span_cache_refcount_underflow():
+    """Test that refcount handles underflow gracefully."""
+    cache = SpanCache()
+
+    # Call stop without start
+    cache.stop()
+
+    # Should work normally after
+    cache.start()
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()

braintrust/test_trace.py

@@ -0,0 +1,267 @@
+"""Tests for Trace functionality."""
+
+import pytest
+from braintrust.trace import CachedSpanFetcher, SpanData
+
+
+# Helper to create mock spans
+def make_span(span_id: str, span_type: str, **extra) -> SpanData:
+    return SpanData(
+        span_id=span_id,
+        input={"text": f"input-{span_id}"},
+        output={"text": f"output-{span_id}"},
+        span_attributes={"type": span_type},
+        **extra,
+    )
+
+
+class TestCachedSpanFetcher:
+    """Test CachedSpanFetcher caching behavior."""
+
+    @pytest.mark.asyncio
+    async def test_fetch_all_spans_without_filter(self):
+        """Test fetching all spans when no filter specified."""
+        mock_spans = [
+            make_span("span-1", "llm"),
+            make_span("span-2", "function"),
+            make_span("span-3", "llm"),
+        ]
+
+        call_count = 0
+
+        async def fetch_fn(span_type):
+            nonlocal call_count
+            call_count += 1
+            return mock_spans
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+        result = await fetcher.get_spans()
+
+        assert call_count == 1
+        assert len(result) == 3
+        assert {s.span_id for s in result} == {"span-1", "span-2", "span-3"}
+
+    @pytest.mark.asyncio
+    async def test_fetch_specific_span_types(self):
+        """Test fetching specific span types when filter specified."""
+        llm_spans = [make_span("span-1", "llm"), make_span("span-2", "llm")]
+
+        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):
+            nonlocal call_count
+            call_count += 1
+            return llm_spans
+
+        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 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.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
+
+    @pytest.mark.asyncio
+    async def test_no_refetch_after_fetching_all_spans(self):
+        """Test that no re-fetching occurs after fetching all spans."""
+        all_spans = [
+            make_span("span-1", "llm"),
+            make_span("span-2", "function"),
+            make_span("span-3", "tool"),
+        ]
+
+        call_count = 0
+
+        async def fetch_fn(span_type):
+            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
+
+        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"}
+
+    @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 fetch_fn(span_type):
+            return all_spans
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+
+        # Fetch all first
+        await fetcher.get_spans()
+
+        # Query for non-existent type
+        result = await fetcher.get_spans(span_type=["nonexistent"])
+        assert len(result) == 0
+
+    @pytest.mark.asyncio
+    async def test_handle_spans_with_no_type(self):
+        """Test handling spans without type (empty string type)."""
+        spans = [
+            make_span("span-1", "llm"),
+            SpanData(span_id="span-2", input={}, span_attributes={}),  # No type
+            SpanData(span_id="span-3", input={}),  # No span_attributes
+        ]
+
+        async def fetch_fn(span_type):
+            return spans
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+
+        # Fetch all
+        result = await fetcher.get_spans()
+        assert len(result) == 3
+
+        # Spans without type go into "" bucket
+        no_type_result = await fetcher.get_spans(span_type=[""])
+        assert len(no_type_result) == 2
+
+    @pytest.mark.asyncio
+    async def test_handle_empty_results(self):
+        """Test handling empty results."""
+
+        async def fetch_fn(span_type):
+            return []
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+
+        result = await fetcher.get_spans()
+        assert len(result) == 0
+
+        # Should still mark as fetched
+        await fetcher.get_spans(span_type=["llm"])
+        # No additional assertions, just making sure it doesn't crash
+
+    @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")]
+
+        call_args = []
+
+        async def fetch_fn(span_type):
+            call_args.append(span_type)
+            return mock_spans
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+
+        result = await fetcher.get_spans(span_type=[])
+
+        assert call_args[0] is None or call_args[0] == []
+        assert len(result) == 1