braintrust 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

braintrust/framework2.py

@@ -1,6 +1,7 @@
 import dataclasses
 import json
-from typing import Any, Callable, Dict, List, Optional, Union, overload
+from collections.abc import Callable
+from typing import Any, overload
 
 import slugify
 from braintrust.logger import api_conn, app_conn, login
@@ -20,7 +21,7 @@ from .util import eprint
 
 class ProjectIdCache:
     def __init__(self):
-        self._cache: Dict[Project, str] = {}
+        self._cache: dict[Project, str] = {}
 
     def get(self, project: "Project") -> str:
         if project not in self._cache:
@@ -31,8 +32,8 @@ class ProjectIdCache:
 
 class _GlobalState:
     def __init__(self):
-        self.functions: List[CodeFunction] = []
-        self.prompts: List[CodePrompt] = []
+        self.functions: list[CodeFunction] = []
+        self.prompts: list[CodePrompt] = []
 
 
 global_ = _GlobalState()
@@ -47,11 +48,11 @@ class CodeFunction:
     name: str
     slug: str
     type_: str
-    description: Optional[str]
+    description: str | None
     parameters: Any
     returns: Any
-    if_exists: Optional[IfExists]
-    metadata: Optional[Dict[str, Any]] = None
+    if_exists: IfExists | None
+    metadata: dict[str, Any] | None = None
 
 
 @dataclasses.dataclass
@@ -62,17 +63,17 @@ class CodePrompt:
     name: str
     slug: str
     prompt: PromptData
-    tool_functions: List[Union[CodeFunction, SavedFunctionId]]
-    description: Optional[str]
-    function_type: Optional[str]
-    id: Optional[str]
-    if_exists: Optional[IfExists]
-    metadata: Optional[Dict[str, Any]] = None
-
-    def to_function_definition(self, if_exists: Optional[IfExists], project_ids: ProjectIdCache) -> Dict[str, Any]:
+    tool_functions: list[CodeFunction | SavedFunctionId]
+    description: str | None
+    function_type: str | None
+    id: str | None
+    if_exists: IfExists | None
+    metadata: dict[str, Any] | None = None
+
+    def to_function_definition(self, if_exists: IfExists | None, project_ids: ProjectIdCache) -> dict[str, Any]:
         prompt_data = self.prompt
         if len(self.tool_functions) > 0:
-            resolvable_tool_functions: List[Any] = []
+            resolvable_tool_functions: list[Any] = []
             for f in self.tool_functions:
                 if isinstance(f, CodeFunction):
                     resolvable_tool_functions.append(
@@ -85,7 +86,7 @@ class CodePrompt:
                 else:
                     resolvable_tool_functions.append(f)
             prompt_data["tool_functions"] = resolvable_tool_functions
-        j: Dict[str, Any] = {
+        j: dict[str, Any] = {
             "project_id": project_ids.get(self.project),
             "name": self.name,
             "slug": self.slug,
@@ -116,13 +117,13 @@ class ToolBuilder:
         self,
         *,
         handler: Callable[..., Any],
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
         parameters: Any = None,
         returns: Any = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> CodeFunction:
         """Creates a tool.
 
@@ -175,48 +176,48 @@ class PromptBuilder:
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        id: Optional[str] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        id: str | None = None,
         prompt: str,
         model: str,
-        params: Optional[ModelParams] = None,
-        tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        params: ModelParams | None = None,
+        tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> CodePrompt: ...
 
     @overload  # messages only, no prompt
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        id: Optional[str] = None,
-        messages: List[ChatCompletionMessageParam],
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        id: str | None = None,
+        messages: list[ChatCompletionMessageParam],
         model: str,
-        params: Optional[ModelParams] = None,
-        tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        params: ModelParams | None = None,
+        tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> CodePrompt: ...
 
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        id: Optional[str] = None,
-        prompt: Optional[str] = None,
-        messages: Optional[List[ChatCompletionMessageParam]] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        id: str | None = None,
+        prompt: str | None = None,
+        messages: list[ChatCompletionMessageParam] | None = None,
         model: str,
-        params: Optional[ModelParams] = None,
-        tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        params: ModelParams | None = None,
+        tools: list[CodeFunction | SavedFunctionId | ToolFunctionDefinition] | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
     ):
         """Creates a prompt.
 
@@ -239,8 +240,8 @@ class PromptBuilder:
         if not slug:
             slug = slugify.slugify(name)
 
-        tool_functions: List[Union[CodeFunction, SavedFunctionId]] = []
-        raw_tools: List[ToolFunctionDefinition] = []
+        tool_functions: list[CodeFunction | SavedFunctionId] = []
+        raw_tools: list[ToolFunctionDefinition] = []
         for tool in tools or []:
             if isinstance(tool, CodeFunction):
                 tool_functions.append(tool)
@@ -298,11 +299,11 @@ class ScorerBuilder:
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
         handler: Callable[..., Any],
         parameters: Any,
         returns: Any = None,
@@ -313,16 +314,16 @@ class ScorerBuilder:
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
         prompt: str,
         model: str,
-        params: Optional[ModelParams] = None,
+        params: ModelParams | None = None,
         use_cot: bool,
-        choice_scores: Dict[str, float],
+        choice_scores: dict[str, float],
     ) -> CodePrompt: ...
 
     # LLM scorer with messages.
@@ -330,38 +331,38 @@ class ScorerBuilder:
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-        messages: List[ChatCompletionMessageParam],
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
+        messages: list[ChatCompletionMessageParam],
         model: str,
-        params: Optional[ModelParams] = None,
+        params: ModelParams | None = None,
         use_cot: bool,
-        choice_scores: Dict[str, float],
+        choice_scores: dict[str, float],
     ) -> CodePrompt: ...
 
     def create(
         self,
         *,
-        name: Optional[str] = None,
-        slug: Optional[str] = None,
-        description: Optional[str] = None,
-        if_exists: Optional[IfExists] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        name: str | None = None,
+        slug: str | None = None,
+        description: str | None = None,
+        if_exists: IfExists | None = None,
+        metadata: dict[str, Any] | None = None,
         # Code scorer params.
-        handler: Optional[Callable[..., Any]] = None,
+        handler: Callable[..., Any] | None = None,
         parameters: Any = None,
         returns: Any = None,
         # LLM scorer params.
-        prompt: Optional[str] = None,
-        messages: Optional[List[ChatCompletionMessageParam]] = None,
-        model: Optional[str] = None,
-        params: Optional[ModelParams] = None,
-        use_cot: Optional[bool] = None,
-        choice_scores: Optional[Dict[str, float]] = None,
-    ) -> Union[CodeFunction, CodePrompt]:
+        prompt: str | None = None,
+        messages: list[ChatCompletionMessageParam] | None = None,
+        model: str | None = None,
+        params: ModelParams | None = None,
+        use_cot: bool | None = None,
+        choice_scores: dict[str, float] | None = None,
+    ) -> CodeFunction | CodePrompt:
         """Creates a scorer.
 
         Args:
@@ -462,8 +463,8 @@ class Project:
         self.prompts = PromptBuilder(self)
         self.scorers = ScorerBuilder(self)
 
-        self._publishable_code_functions: List[CodeFunction] = []
-        self._publishable_prompts: List[CodePrompt] = []
+        self._publishable_code_functions: list[CodeFunction] = []
+        self._publishable_prompts: list[CodePrompt] = []
 
     def add_code_function(self, fn: CodeFunction):
         self._publishable_code_functions.append(fn)
@@ -483,7 +484,7 @@ class Project:
         login()
         project_id_cache = ProjectIdCache()
 
-        definitions: List[Dict[str, Any]] = []
+        definitions: list[dict[str, Any]] = []
         if self._publishable_code_functions:
             eprint(
                 f"{bcolors.WARNING}Code functions cannot be published directly. Use `braintrust push` instead.{bcolors.ENDC}"

braintrust/functions/invoke.py

@@ -1,38 +1,58 @@
-from typing import Any, Dict, List, Literal, Optional, TypeVar, Union, overload
+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 ..util import response_raise_for_status
 from .constants import INVOKE_API_VERSION
 from .stream import BraintrustInvokeError, BraintrustStream
 
 T = TypeVar("T")
-ModeType = Literal["auto", "parallel"]
+ModeType = Literal["auto", "parallel", "json", "text"]
+ObjectType = Literal["project_logs", "experiment", "dataset", "playground_logs"]
+
+
+class SpanScope(TypedDict):
+    """Scope for operating on a single span."""
+
+    type: Literal["span"]
+    id: str
+    root_span_id: str
+
+
+class TraceScope(TypedDict):
+    """Scope for operating on an entire trace."""
+
+    type: Literal["trace"]
+    root_span_id: str
 
 
 @overload
 def invoke(
     # the permutations of arguments for a function id
-    function_id: Optional[str] = None,
-    version: Optional[str] = None,
-    prompt_session_id: Optional[str] = None,
-    prompt_session_function_id: Optional[str] = None,
-    project_name: Optional[str] = None,
-    slug: Optional[str] = None,
-    global_function: Optional[str] = None,
+    function_id: str | None = None,
+    version: str | None = None,
+    prompt_session_id: str | None = None,
+    prompt_session_function_id: str | None = None,
+    project_name: str | None = None,
+    project_id: str | None = None,
+    slug: str | None = None,
+    global_function: str | None = None,
+    function_type: FunctionTypeEnum | None = None,
     # arguments to the function
     input: Any = None,
-    messages: Optional[List[Any]] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-    tags: Optional[List[str]] = None,
-    parent: Optional[Union[Exportable, str]] = None,
-    stream: Optional[Literal[False]] = None,
-    mode: Optional[ModeType] = None,
-    strict: Optional[bool] = None,
-    org_name: Optional[str] = None,
-    api_key: Optional[str] = None,
-    app_url: Optional[str] = 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,
+    stream: Literal[False] | None = None,
+    mode: ModeType | None = None,
+    strict: bool | None = None,
+    org_name: str | None = None,
+    api_key: str | None = None,
+    app_url: str | None = None,
     force_login: bool = False,
 ) -> T: ...
 
@@ -40,52 +60,58 @@ def invoke(
 @overload
 def invoke(
     # the permutations of arguments for a function id
-    function_id: Optional[str] = None,
-    version: Optional[str] = None,
-    prompt_session_id: Optional[str] = None,
-    prompt_session_function_id: Optional[str] = None,
-    project_name: Optional[str] = None,
-    slug: Optional[str] = None,
-    global_function: Optional[str] = None,
+    function_id: str | None = None,
+    version: str | None = None,
+    prompt_session_id: str | None = None,
+    prompt_session_function_id: str | None = None,
+    project_name: str | None = None,
+    project_id: str | None = None,
+    slug: str | None = None,
+    global_function: str | None = None,
+    function_type: FunctionTypeEnum | None = None,
     # arguments to the function
     input: Any = None,
-    messages: Optional[List[Any]] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-    tags: Optional[List[str]] = None,
-    parent: Optional[Union[Exportable, str]] = 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,
     stream: Literal[True] = True,
-    mode: Optional[ModeType] = None,
-    strict: Optional[bool] = None,
-    org_name: Optional[str] = None,
-    api_key: Optional[str] = None,
-    app_url: Optional[str] = None,
+    mode: ModeType | None = None,
+    strict: bool | None = None,
+    org_name: str | None = None,
+    api_key: str | None = None,
+    app_url: str | None = None,
     force_login: bool = False,
 ) -> BraintrustStream: ...
 
 
 def invoke(
     # the permutations of arguments for a function id
-    function_id: Optional[str] = None,
-    version: Optional[str] = None,
-    prompt_session_id: Optional[str] = None,
-    prompt_session_function_id: Optional[str] = None,
-    project_name: Optional[str] = None,
-    slug: Optional[str] = None,
-    global_function: Optional[str] = None,
+    function_id: str | None = None,
+    version: str | None = None,
+    prompt_session_id: str | None = None,
+    prompt_session_function_id: str | None = None,
+    project_name: str | None = None,
+    project_id: str | None = None,
+    slug: str | None = None,
+    global_function: str | None = None,
+    function_type: FunctionTypeEnum | None = None,
     # arguments to the function
     input: Any = None,
-    messages: Optional[List[Any]] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-    tags: Optional[List[str]] = None,
-    parent: Optional[Union[Exportable, str]] = 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,
     stream: bool = False,
-    mode: Optional[ModeType] = None,
-    strict: Optional[bool] = None,
-    org_name: Optional[str] = None,
-    api_key: Optional[str] = None,
-    app_url: Optional[str] = None,
+    mode: ModeType | None = None,
+    strict: bool | None = None,
+    org_name: str | None = None,
+    api_key: str | None = None,
+    app_url: str | None = None,
     force_login: bool = False,
-) -> Union[BraintrustStream, T]:
+) -> BraintrustStream | T:
     """
     Invoke a Braintrust function, returning a `BraintrustStream` or the value as a plain
     Python object.
@@ -93,6 +119,8 @@ 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.
@@ -118,8 +146,12 @@ def invoke(
         prompt_session_id: The ID of the prompt session to invoke the function from.
         prompt_session_function_id: The ID of the function in the prompt session to invoke.
         project_name: The name of the project containing the function to invoke.
+        project_id: The ID of the project to use for execution context (API keys, project defaults, etc.).
+            This is not the project the function belongs to, but the project context for the invocation.
         slug: The slug of the function to invoke.
         global_function: The name of the global function to invoke.
+        function_type: The type of the global function to invoke. If unspecified, defaults to 'scorer'
+            for backward compatibility.
 
     Returns:
         The output of the function. If `stream` is True, returns a `BraintrustStream`,
@@ -149,6 +181,8 @@ def invoke(
         function_id_args["slug"] = slug
     if global_function is not None:
         function_id_args["global_function"] = global_function
+    if function_type is not None:
+        function_id_args["function_type"] = function_type
 
     request = dict(
         input=input,
@@ -161,12 +195,18 @@ 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:
         request["strict"] = strict
 
     headers = {"Accept": "text/event-stream" if stream else "application/json"}
+    if project_id is not None:
+        headers["x-bt-project-id"] = project_id
+    if org_name is not None:
+        headers["x-bt-org-name"] = org_name
 
     resp = proxy_conn().post("function/invoke", json=request, headers=headers, stream=stream)
     if resp.status_code == 500:
@@ -180,7 +220,7 @@ def invoke(
         return resp.json()
 
 
-def init_function(project_name: str, slug: str, version: Optional[str] = None):
+def init_function(project_name: str, slug: str, version: str | None = None):
     """
     Creates a function that can be used as either a task or scorer in the Eval framework.
     When used as a task, it will invoke the specified Braintrust function with the input.

braintrust/functions/stream.py

@@ -7,8 +7,9 @@ with utility methods to make them easy to log and convert into various formats.
 
 import dataclasses
 import json
+from collections.abc import Generator
 from itertools import tee
-from typing import Generator, List, Literal, Union
+from typing import Literal, Union
 
 from sseclient import SSEClient
 
@@ -87,7 +88,7 @@ class BraintrustStream:
     with utility methods to make them easy to log and convert into various formats.
     """
 
-    def __init__(self, base_stream: Union[SSEClient, List[BraintrustStreamChunk]]):
+    def __init__(self, base_stream: SSEClient | list[BraintrustStreamChunk]):
         """
         Initialize a BraintrustStream.
 

braintrust/generated_types.py

@@ -1,4 +1,4 @@
-"""Auto-generated file (internal git SHA 8e9c0a96b3cf291360978c17580f72f6817bd6c8) -- do not modify"""
+"""Auto-generated file (internal git SHA 547fa17c0937e0e25fdf9214487be6f31c91a37a) -- do not modify"""
 
 from ._generated_types import (
     Acl,
@@ -32,6 +32,7 @@ from ._generated_types import (
     ExperimentEvent,
     ExtendedSavedFunctionId,
     ExternalAttachmentReference,
+    FacetData,
     Function,
     FunctionData,
     FunctionFormat,
@@ -47,10 +48,15 @@ from ._generated_types import (
     GraphNode,
     Group,
     IfExists,
+    InvokeContext,
     InvokeFunction,
     InvokeParent,
+    InvokeScope,
+    MCPServer,
     MessageRole,
     ModelParams,
+    NullableFunctionTypeEnum,
+    NullableSavedFunctionId,
     ObjectReference,
     ObjectReferenceNullish,
     OnlineScoreConfig,
@@ -86,11 +92,13 @@ from ._generated_types import (
     ServiceToken,
     SpanAttributes,
     SpanIFrame,
+    SpanScope,
     SpanType,
     SSEConsoleEventData,
     SSEProgressEventData,
     StreamingMode,
     ToolFunctionDefinition,
+    TraceScope,
     UploadStatus,
     User,
     View,
@@ -131,6 +139,7 @@ __all__ = [
     "ExperimentEvent",
     "ExtendedSavedFunctionId",
     "ExternalAttachmentReference",
+    "FacetData",
     "Function",
     "FunctionData",
     "FunctionFormat",
@@ -146,10 +155,15 @@ __all__ = [
     "GraphNode",
     "Group",
     "IfExists",
+    "InvokeContext",
     "InvokeFunction",
     "InvokeParent",
+    "InvokeScope",
+    "MCPServer",
     "MessageRole",
     "ModelParams",
+    "NullableFunctionTypeEnum",
+    "NullableSavedFunctionId",
     "ObjectReference",
     "ObjectReferenceNullish",
     "OnlineScoreConfig",
@@ -187,9 +201,11 @@ __all__ = [
     "ServiceToken",
     "SpanAttributes",
     "SpanIFrame",
+    "SpanScope",
     "SpanType",
     "StreamingMode",
     "ToolFunctionDefinition",
+    "TraceScope",
     "UploadStatus",
     "User",
     "View",

braintrust/git_fields.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import List, Literal, Optional
+from typing import Literal
 
 from .serializable_data_class import SerializableDataClass
 
@@ -8,21 +8,21 @@ from .serializable_data_class import SerializableDataClass
 class RepoInfo(SerializableDataClass):
     """Information about the current HEAD of the repo."""
 
-    commit: Optional[str] = None
-    branch: Optional[str] = None
-    tag: Optional[str] = None
-    dirty: Optional[bool] = None
-    author_name: Optional[str] = None
-    author_email: Optional[str] = None
-    commit_message: Optional[str] = None
-    commit_time: Optional[str] = None
-    git_diff: Optional[str] = None
+    commit: str | None = None
+    branch: str | None = None
+    tag: str | None = None
+    dirty: bool | None = None
+    author_name: str | None = None
+    author_email: str | None = None
+    commit_message: str | None = None
+    commit_time: str | None = None
+    git_diff: str | None = None
 
 
 @dataclass
 class GitMetadataSettings(SerializableDataClass):
     collect: Literal["all", "some", "none"] = "all"
-    fields: Optional[List[str]] = field(default_factory=list)
+    fields: list[str] | None = field(default_factory=list)
 
     @classmethod
     def merge(cls, s1: "GitMetadataSettings", s2: "GitMetadataSettings") -> "GitMetadataSettings":