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

braintrust/cli/install/api.py

@@ -326,7 +326,7 @@ def main(args):
             textwrap.dedent(
                 f"""\
             Stack with name {args.name} does not exist. Either create it manually by following
-            https://www.braintrust.dev/docs/guides/self-hosting/aws or use the --create flag."""
+            https://www.braintrust.dev/docs/admin/self-hosting/aws or use the --create flag."""
             )
         )
         exit(1)

braintrust/conftest.py

@@ -46,3 +46,27 @@ def reset_braintrust_state():
     from braintrust import logger
 
     logger._state = logger.BraintrustState()
+
+
+@pytest.fixture(scope="session")
+def vcr_config():
+    """
+    VCR configuration for recording/playing back HTTP interactions.
+
+    In CI, use "none" to fail if cassette is missing.
+    Locally, use "once" to record new cassettes if they don't exist.
+    """
+    record_mode = "none" if (os.environ.get("CI") or os.environ.get("GITHUB_ACTIONS")) else "once"
+
+    return {
+        "record_mode": record_mode,
+        "filter_headers": [
+            "authorization",
+            "openai-organization",
+            "x-api-key",
+            "api-key",
+            "openai-api-key",
+            "x-goog-api-key",
+            "x-bt-auth-token",
+        ],
+    }

braintrust/devserver/test_server_integration.py

@@ -8,17 +8,6 @@ from braintrust.framework import _evals
 from braintrust.test_helpers import has_devserver_installed
 
 
-@pytest.fixture(scope="module")
-def vcr_config():
-    """VCR configuration to filter sensitive headers."""
-    return {
-        "filter_headers": [
-            "x-bt-auth-token",
-            "authorization",
-        ]
-    }
-
-
 @pytest.fixture
 def client():
     """Create test client using the real simple_eval.py example."""

braintrust/framework.py

@@ -1280,6 +1280,29 @@ async def _run_evaluator_internal(
     filters: list[Filter],
     stream: Callable[[SSEProgressEvent], None] | None = None,
     state: BraintrustState | None = None,
+):
+    # Start span cache for this eval (it's disabled by default to avoid temp files outside of evals)
+    if state is None:
+        from braintrust.logger import _internal_get_global_state
+
+        state = _internal_get_global_state()
+
+    state.span_cache.start()
+    try:
+        return await _run_evaluator_internal_impl(experiment, evaluator, position, filters, stream, state)
+    finally:
+        # Clean up disk-based span cache after eval completes and stop caching
+        state.span_cache.dispose()
+        state.span_cache.stop()
+
+
+async def _run_evaluator_internal_impl(
+    experiment,
+    evaluator: Evaluator,
+    position: int | None,
+    filters: list[Filter],
+    stream: Callable[[SSEProgressEvent], None] | None = None,
+    state: BraintrustState | None = None,
 ):
     event_loop = asyncio.get_event_loop()
 
@@ -1290,11 +1313,13 @@ async def _run_evaluator_internal(
             {**parent_propagated},
             {"span_attributes": {"purpose": "scorer"}},
         )
+        # Strip trace from logged input - it's internal plumbing that shouldn't appear in spans
+        logged_input = {k: v for k, v in kwargs.items() if k != "trace"}
         with root_span.start_span(
             name=name,
             span_attributes={"type": SpanTypeAttribute.SCORE, "purpose": "scorer"},
             propagated_event=merged_propagated,
-            input=dict(**kwargs),
+            input=logged_input,
         ) as span:
             score = scorer
             if hasattr(scorer, "eval_async"):
@@ -1415,6 +1440,77 @@ async def _run_evaluator_internal(
                 tags = hooks.tags if hooks.tags else None
                 root_span.log(output=output, metadata=metadata, tags=tags)
 
+                # Create trace object for scorers
+                from braintrust.trace import LocalTrace
+
+                async def ensure_spans_flushed():
+                    # Flush native Braintrust spans
+                    if experiment:
+                        await asyncio.get_event_loop().run_in_executor(
+                            None, lambda: experiment.state.flush()
+                        )
+                    elif state:
+                        await asyncio.get_event_loop().run_in_executor(None, lambda: state.flush())
+                    else:
+                        from braintrust.logger import flush as flush_logger
+
+                        await asyncio.get_event_loop().run_in_executor(None, flush_logger)
+
+                    # Also flush OTEL spans if registered
+                    if state:
+                        await state.flush_otel()
+
+                experiment_id = None
+                if experiment:
+                    try:
+                        experiment_id = experiment.id
+                    except:
+                        experiment_id = None
+
+                trace = None
+                if state or experiment:
+                    # Get the state to use
+                    trace_state = state
+                    if not trace_state and experiment:
+                        trace_state = experiment.state
+                    if not trace_state:
+                        # Fall back to global state
+                        from braintrust.logger import _internal_get_global_state
+
+                        trace_state = _internal_get_global_state()
+
+                    # Access root_span_id from the concrete SpanImpl instance
+                    # The Span interface doesn't expose this but SpanImpl has it
+                    root_span_id_value = getattr(root_span, "root_span_id", root_span.id)
+
+                    # Check if there's a parent in the context to determine object_type and object_id
+                    from braintrust.span_identifier_v3 import SpanComponentsV3, span_object_type_v3_to_typed_string
+
+                    parent_str = trace_state.current_parent.get()
+                    parent_components = None
+                    if parent_str:
+                        try:
+                            parent_components = SpanComponentsV3.from_str(parent_str)
+                        except Exception:
+                            # If parsing fails, parent_components stays None
+                            pass
+
+                    # Determine object_type and object_id based on parent or experiment
+                    if parent_components:
+                        trace_object_type = span_object_type_v3_to_typed_string(parent_components.object_type)
+                        trace_object_id = parent_components.object_id or ""
+                    else:
+                        trace_object_type = "experiment"
+                        trace_object_id = experiment_id or ""
+
+                    trace = LocalTrace(
+                        object_type=trace_object_type,
+                        object_id=trace_object_id,
+                        root_span_id=root_span_id_value,
+                        ensure_spans_flushed=ensure_spans_flushed,
+                        state=trace_state,
+                    )
+
                 score_promises = [
                     asyncio.create_task(
                         await_or_run_scorer(
@@ -1426,6 +1522,7 @@ async def _run_evaluator_internal(
                                 "expected": datum.expected,
                                 "metadata": metadata,
                                 "output": output,
+                                "trace": trace,
                             },
                         )
                     )

braintrust/functions/invoke.py

@@ -2,8 +2,8 @@ from typing import Any, Literal, TypedDict, TypeVar, overload
 
 from sseclient import SSEClient
 
-from .._generated_types import FunctionTypeEnum, InvokeContext
-from ..logger import Exportable, get_span_parent_object, login, proxy_conn
+from .._generated_types import FunctionTypeEnum
+from ..logger import Exportable, _internal_get_global_state, get_span_parent_object, login, proxy_conn
 from ..util import response_raise_for_status
 from .constants import INVOKE_API_VERSION
 from .stream import BraintrustInvokeError, BraintrustStream
@@ -43,7 +43,6 @@ def invoke(
     # arguments to the function
     input: Any = None,
     messages: list[Any] | None = None,
-    context: InvokeContext | None = None,
     metadata: dict[str, Any] | None = None,
     tags: list[str] | None = None,
     parent: Exportable | str | None = None,
@@ -72,7 +71,6 @@ def invoke(
     # arguments to the function
     input: Any = None,
     messages: list[Any] | None = None,
-    context: InvokeContext | None = None,
     metadata: dict[str, Any] | None = None,
     tags: list[str] | None = None,
     parent: Exportable | str | None = None,
@@ -100,7 +98,6 @@ def invoke(
     # arguments to the function
     input: Any = None,
     messages: list[Any] | None = None,
-    context: InvokeContext | None = None,
     metadata: dict[str, Any] | None = None,
     tags: list[str] | None = None,
     parent: Exportable | str | None = None,
@@ -119,8 +116,6 @@ def invoke(
     Args:
         input: The input to the function. This will be logged as the `input` field in the span.
         messages: Additional OpenAI-style messages to add to the prompt (only works for llm functions).
-        context: Context for functions that operate on spans/traces (e.g., facets). Should contain
-            `object_type`, `object_id`, and `scope` fields.
         metadata: Additional metadata to add to the span. This will be logged as the `metadata` field in the span.
             It will also be available as the {{metadata}} field in the prompt and as the `metadata` argument
             to the function.
@@ -195,8 +190,6 @@ def invoke(
     )
     if messages is not None:
         request["messages"] = messages
-    if context is not None:
-        request["context"] = context
     if mode is not None:
         request["mode"] = mode
     if strict is not None:
@@ -250,6 +243,8 @@ def init_function(project_name: str, slug: str, version: str | None = None):
     :param version: Optional version of the function to use. Defaults to latest.
     :return: A function that can be used as a task or scorer.
     """
+    # Disable span cache since remote function spans won't be in the local cache
+    _internal_get_global_state().span_cache.disable()
 
     def f(*args: Any, **kwargs: Any) -> Any:
         if len(args) > 0:

braintrust/functions/test_invoke.py

@@ -0,0 +1,61 @@
+"""Tests for the invoke module, particularly init_function."""
+
+
+from braintrust.functions.invoke import init_function
+from braintrust.logger import _internal_get_global_state, _internal_reset_global_state
+
+
+class TestInitFunction:
+    """Tests for init_function."""
+
+    def setup_method(self):
+        """Reset state before each test."""
+        _internal_reset_global_state()
+
+    def teardown_method(self):
+        """Clean up after each test."""
+        _internal_reset_global_state()
+
+    def test_init_function_disables_span_cache(self):
+        """Test that init_function disables the span cache."""
+        state = _internal_get_global_state()
+
+        # Cache should be disabled by default (it's only enabled during evals)
+        assert state.span_cache.disabled is True
+
+        # Enable the cache (simulating what happens during eval)
+        state.span_cache.start()
+        assert state.span_cache.disabled is False
+
+        # Call init_function
+        f = init_function("test-project", "test-function")
+
+        # Cache should now be disabled (init_function explicitly disables it)
+        assert state.span_cache.disabled is True
+        assert f.__name__ == "init_function-test-project-test-function-latest"
+
+    def test_init_function_with_version(self):
+        """Test that init_function creates a function with the correct name including version."""
+        f = init_function("my-project", "my-scorer", version="v1")
+        assert f.__name__ == "init_function-my-project-my-scorer-v1"
+
+    def test_init_function_without_version_uses_latest(self):
+        """Test that init_function uses 'latest' in name when version not specified."""
+        f = init_function("my-project", "my-scorer")
+        assert f.__name__ == "init_function-my-project-my-scorer-latest"
+
+    def test_init_function_permanently_disables_cache(self):
+        """Test that init_function permanently disables the cache (can't be re-enabled)."""
+        state = _internal_get_global_state()
+
+        # Enable the cache
+        state.span_cache.start()
+        assert state.span_cache.disabled is False
+
+        # Call init_function
+        init_function("test-project", "test-function")
+        assert state.span_cache.disabled is True
+
+        # Try to start again - should still be disabled because of explicit disable
+        state.span_cache.start()
+        assert state.span_cache.disabled is True

braintrust/generated_types.py

@@ -1,4 +1,4 @@
-"""Auto-generated file (internal git SHA 547fa17c0937e0e25fdf9214487be6f31c91a37a) -- do not modify"""
+"""Auto-generated file (internal git SHA 21146f64bf5ad1eadd3a99d186274728e25e5399) -- do not modify"""
 
 from ._generated_types import (
     Acl,
@@ -10,6 +10,7 @@ from ._generated_types import (
     AsyncScoringState,
     AttachmentReference,
     AttachmentStatus,
+    BatchedFacetData,
     BraintrustAttachmentReference,
     BraintrustModelParams,
     CallEvent,
@@ -28,6 +29,9 @@ from ._generated_types import (
     Dataset,
     DatasetEvent,
     EnvVar,
+    EvalStatusPage,
+    EvalStatusPageConfig,
+    EvalStatusPageTheme,
     Experiment,
     ExperimentEvent,
     ExtendedSavedFunctionId,
@@ -47,15 +51,13 @@ from ._generated_types import (
     GraphEdge,
     GraphNode,
     Group,
+    GroupScope,
     IfExists,
-    InvokeContext,
     InvokeFunction,
     InvokeParent,
-    InvokeScope,
     MCPServer,
     MessageRole,
     ModelParams,
-    NullableFunctionTypeEnum,
     NullableSavedFunctionId,
     ObjectReference,
     ObjectReferenceNullish,
@@ -99,6 +101,7 @@ from ._generated_types import (
     StreamingMode,
     ToolFunctionDefinition,
     TraceScope,
+    TriggeredFunctionState,
     UploadStatus,
     User,
     View,
@@ -117,6 +120,7 @@ __all__ = [
     "AsyncScoringState",
     "AttachmentReference",
     "AttachmentStatus",
+    "BatchedFacetData",
     "BraintrustAttachmentReference",
     "BraintrustModelParams",
     "CallEvent",
@@ -135,6 +139,9 @@ __all__ = [
     "Dataset",
     "DatasetEvent",
     "EnvVar",
+    "EvalStatusPage",
+    "EvalStatusPageConfig",
+    "EvalStatusPageTheme",
     "Experiment",
     "ExperimentEvent",
     "ExtendedSavedFunctionId",
@@ -154,15 +161,13 @@ __all__ = [
     "GraphEdge",
     "GraphNode",
     "Group",
+    "GroupScope",
     "IfExists",
-    "InvokeContext",
     "InvokeFunction",
     "InvokeParent",
-    "InvokeScope",
     "MCPServer",
     "MessageRole",
     "ModelParams",
-    "NullableFunctionTypeEnum",
     "NullableSavedFunctionId",
     "ObjectReference",
     "ObjectReferenceNullish",
@@ -206,6 +211,7 @@ __all__ = [
     "StreamingMode",
     "ToolFunctionDefinition",
     "TraceScope",
+    "TriggeredFunctionState",
     "UploadStatus",
     "User",
     "View",

braintrust/logger.py

@@ -47,12 +47,9 @@ from urllib3.util.retry import Retry
 from . import context, id_gen
 from .bt_json import bt_dumps, bt_safe_deep_copy
 from .db_fields import (
-    ASYNC_SCORING_CONTROL_FIELD,
     AUDIT_METADATA_FIELD,
     AUDIT_SOURCE_FIELD,
     IS_MERGE_FIELD,
-    MERGE_PATHS_FIELD,
-    SKIP_ASYNC_SCORING_FIELD,
     TRANSACTION_ID_FIELD,
     VALID_SOURCES,
 )
@@ -101,6 +98,14 @@ from .xact_ids import prettify_xact
 Metadata = dict[str, Any]
 DATA_API_VERSION = 2
 
+
+class DatasetRef(TypedDict, total=False):
+    """Reference to a dataset by ID and optional version."""
+
+    id: str
+    version: str
+
+
 T = TypeVar("T")
 TMapping = TypeVar("TMapping", bound=Mapping[str, Any])
 TMutableMapping = TypeVar("TMutableMapping", bound=MutableMapping[str, Any])
@@ -396,6 +401,11 @@ class BraintrustState:
             ),
         )
 
+        from braintrust.span_cache import SpanCache
+
+        self.span_cache = SpanCache()
+        self._otel_flush_callback: Any | None = None
+
     def reset_login_info(self):
         self.app_url: str | None = None
         self.app_public_url: str | None = None
@@ -452,26 +462,39 @@ class BraintrustState:
 
         return self._context_manager
 
+    def register_otel_flush(self, callback: Any) -> None:
+        """
+        Register an OTEL flush callback. This is called by the OTEL integration
+        when it initializes a span processor/exporter.
+        """
+        self._otel_flush_callback = callback
+
+    async def flush_otel(self) -> None:
+        """
+        Flush OTEL spans if a callback is registered.
+        Called during ensure_spans_flushed to ensure OTEL spans are visible in BTQL.
+        """
+        if self._otel_flush_callback:
+            await self._otel_flush_callback()
+
     def copy_state(self, other: "BraintrustState"):
         """Copy login information from another BraintrustState instance."""
-        self.__dict__.update(
-            {
-                k: v
-                for (k, v) in other.__dict__.items()
-                if k
-                not in (
-                    "current_experiment",
-                    "current_logger",
-                    "current_parent",
-                    "current_span",
-                    "_global_bg_logger",
-                    "_override_bg_logger",
-                    "_context_manager",
-                    "_last_otel_setting",
-                    "_context_manager_lock",
-                )
-            }
-        )
+        self.__dict__.update({
+            k: v
+            for (k, v) in other.__dict__.items()
+            if k
+            not in (
+                "current_experiment",
+                "current_logger",
+                "current_parent",
+                "current_span",
+                "_global_bg_logger",
+                "_override_bg_logger",
+                "_context_manager",
+                "_last_otel_setting",
+                "_context_manager_lock",
+            )
+        })
 
     def login(
         self,
@@ -1299,7 +1322,7 @@ def init(
     project: str | None = None,
     experiment: str | None = None,
     description: str | None = None,
-    dataset: Optional["Dataset"] = None,
+    dataset: Optional["Dataset"] | DatasetRef = None,
     open: bool = False,
     base_experiment: str | None = None,
     is_public: bool = False,
@@ -1412,12 +1435,19 @@ def init(
             args["base_exp_id"] = base_experiment_id
         elif base_experiment is not None:
             args["base_experiment"] = base_experiment
-        else:
+        elif merged_git_metadata_settings and merged_git_metadata_settings.collect != "none":
             args["ancestor_commits"] = list(get_past_n_ancestors())
 
         if dataset is not None:
-            args["dataset_id"] = dataset.id
-            args["dataset_version"] = dataset.version
+            if isinstance(dataset, dict):
+                # Simple {"id": ..., "version": ...} dict
+                args["dataset_id"] = dataset["id"]
+                if "version" in dataset:
+                    args["dataset_version"] = dataset["version"]
+            else:
+                # Full Dataset object
+                args["dataset_id"] = dataset.id
+                args["dataset_version"] = dataset.version
 
         if is_public is not None:
             args["public"] = is_public
@@ -1448,7 +1478,11 @@ def init(
     # For experiments, disable queue size limit enforcement (unlimited queue)
     state.enforce_queue_size_limit(False)
 
-    ret = Experiment(lazy_metadata=LazyValue(compute_metadata, use_mutex=True), dataset=dataset, state=state)
+    ret = Experiment(
+        lazy_metadata=LazyValue(compute_metadata, use_mutex=True),
+        dataset=dataset if isinstance(dataset, Dataset) else None,
+        state=state,
+    )
     if set_current:
         state.current_experiment = ret
     return ret
@@ -1763,6 +1797,25 @@ def login(
         _state.login(app_url=app_url, api_key=api_key, org_name=org_name, force_login=force_login)
 
 
+def register_otel_flush(callback: Any) -> None:
+    """
+    Register a callback to flush OTEL spans. This is called by the OTEL integration
+    when it initializes a span processor/exporter.
+
+    When ensure_spans_flushed is called (e.g., before a BTQL query in scorers),
+    this callback will be invoked to ensure OTEL spans are flushed to the server.
+
+    Also disables the span cache, since OTEL spans aren't in the local cache
+    and we need BTQL to see the complete span tree (both native + OTEL spans).
+
+    :param callback: The async callback function to flush OTEL spans.
+    """
+    global _state
+    _state.register_otel_flush(callback)
+    # Disable span cache since OTEL spans aren't in the local cache
+    _state.span_cache.disable()
+
+
 def login_to_state(
     app_url: str | None = None,
     api_key: str | None = None,
@@ -2325,29 +2378,6 @@ def _enrich_attachments(event: TMutableMapping) -> TMutableMapping:
 
 
 def _validate_and_sanitize_experiment_log_partial_args(event: Mapping[str, Any]) -> dict[str, Any]:
-    # Make sure only certain keys are specified.
-    forbidden_keys = set(event.keys()) - {
-        "input",
-        "output",
-        "expected",
-        "tags",
-        "scores",
-        "metadata",
-        "metrics",
-        "error",
-        "dataset_record_id",
-        "origin",
-        "inputs",
-        "span_attributes",
-        ASYNC_SCORING_CONTROL_FIELD,
-        MERGE_PATHS_FIELD,
-        SKIP_ASYNC_SCORING_FIELD,
-        "span_id",
-        "root_span_id",
-    }
-    if forbidden_keys:
-        raise ValueError(f"The following keys are not permitted: {forbidden_keys}")
-
     scores = event.get("scores")
     if scores:
         for name, score in scores.items():
@@ -3856,6 +3886,21 @@ class SpanImpl(Span):
         if serializable_partial_record.get("metrics", {}).get("end") is not None:
             self._logged_end_time = serializable_partial_record["metrics"]["end"]
 
+        # Write to local span cache for scorer access
+        # Only cache experiment spans - regular logs don't need caching
+        if self.parent_object_type == SpanObjectTypeV3.EXPERIMENT:
+            from braintrust.span_cache import CachedSpan
+
+            cached_span = CachedSpan(
+                span_id=self.span_id,
+                input=serializable_partial_record.get("input"),
+                output=serializable_partial_record.get("output"),
+                metadata=serializable_partial_record.get("metadata"),
+                span_parents=self.span_parents,
+                span_attributes=serializable_partial_record.get("span_attributes"),
+            )
+            self.state.span_cache.queue_write(self.root_span_id, self.span_id, cached_span)
+
         def compute_record() -> dict[str, Any]:
             exporter = _get_exporter()
             return dict(
@@ -4403,24 +4448,20 @@ def render_message(render: Callable[[str], str], message: PromptMessage):
                 if c["type"] == "text":
                     rendered_content.append({**c, "text": render(c["text"])})
                 elif c["type"] == "image_url":
-                    rendered_content.append(
-                        {
-                            **c,
-                            "image_url": {**c["image_url"], "url": render(c["image_url"]["url"])},
-                        }
-                    )
+                    rendered_content.append({
+                        **c,
+                        "image_url": {**c["image_url"], "url": render(c["image_url"]["url"])},
+                    })
                 elif c["type"] == "file":
-                    rendered_content.append(
-                        {
-                            **c,
-                            "file": {
-                                **c["file"],
-                                "file_data": render(c["file"]["file_data"]),
-                                **({} if "file_id" not in c["file"] else {"file_id": render(c["file"]["file_id"])}),
-                                **({} if "filename" not in c["file"] else {"filename": render(c["file"]["filename"])}),
-                            },
-                        }
-                    )
+                    rendered_content.append({
+                        **c,
+                        "file": {
+                            **c["file"],
+                            "file_data": render(c["file"]["file_data"]),
+                            **({} if "file_id" not in c["file"] else {"file_id": render(c["file"]["file_id"])}),
+                            **({} if "filename" not in c["file"] else {"filename": render(c["file"]["filename"])}),
+                        },
+                    })
                 else:
                     raise ValueError(f"Unknown content type: {c['type']}")
 

braintrust/prompt_cache/test_disk_cache.py

@@ -39,7 +39,7 @@ class TestDiskCache(unittest.TestCase):
             "a\nb",
         ]
         for k in weird_keys:
-            time.sleep(0.05)  # make sure the mtimes are different
+            time.sleep(0.01)  # make sure the mtimes are different
             self.cache.set(k, data)
             result = self.cache.get(k)
             assert data == result
@@ -61,7 +61,7 @@ class TestDiskCache(unittest.TestCase):
         # Fill cache beyond max size (3).
         for i in range(3):
             self.cache.set(f"key{i}", {"value": i})
-            time.sleep(0.1)  # wait to ensure different mtimes
+            time.sleep(0.01)  # wait to ensure different mtimes
 
         # Add one more to trigger eviction.
         self.cache.set("key3", {"value": 3})
@@ -75,7 +75,7 @@ class TestDiskCache(unittest.TestCase):
         # Fill cache beyond max size (3).
         for i in range(3):
             self.cache.set(f"key{i}", {"value": i})
-            time.sleep(0.1)  # wait to ensure different mtimes
+            time.sleep(0.01)  # wait to ensure different mtimes
 
         # Add one more to trigger eviction.
         self.cache.set("key3", {"value": 3})