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

braintrust/conftest.py

@@ -36,6 +36,7 @@ def override_app_url_for_tests():
 @pytest.fixture(autouse=True)
 def setup_braintrust():
     os.environ.setdefault("GOOGLE_API_KEY", os.getenv("GEMINI_API_KEY", "your_google_api_key_here"))
+    os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy-api-key-for-vcr-tests")
 
 
 @pytest.fixture(autouse=True)

braintrust/functions/invoke.py

@@ -1,14 +1,31 @@
-from typing import Any, Dict, List, Literal, Optional, TypeVar, Union, overload
+from typing import Any, Dict, List, Literal, Optional, TypedDict, TypeVar, Union, overload
 
 from sseclient import SSEClient
 
+from .._generated_types import 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
@@ -19,11 +36,13 @@ def invoke(
     prompt_session_id: Optional[str] = None,
     prompt_session_function_id: Optional[str] = None,
     project_name: Optional[str] = None,
+    project_id: Optional[str] = None,
     slug: Optional[str] = None,
     global_function: Optional[str] = None,
     # arguments to the function
     input: Any = None,
     messages: Optional[List[Any]] = None,
+    context: Optional[InvokeContext] = None,
     metadata: Optional[Dict[str, Any]] = None,
     tags: Optional[List[str]] = None,
     parent: Optional[Union[Exportable, str]] = None,
@@ -45,11 +64,13 @@ def invoke(
     prompt_session_id: Optional[str] = None,
     prompt_session_function_id: Optional[str] = None,
     project_name: Optional[str] = None,
+    project_id: Optional[str] = None,
     slug: Optional[str] = None,
     global_function: Optional[str] = None,
     # arguments to the function
     input: Any = None,
     messages: Optional[List[Any]] = None,
+    context: Optional[InvokeContext] = None,
     metadata: Optional[Dict[str, Any]] = None,
     tags: Optional[List[str]] = None,
     parent: Optional[Union[Exportable, str]] = None,
@@ -70,11 +91,13 @@ def invoke(
     prompt_session_id: Optional[str] = None,
     prompt_session_function_id: Optional[str] = None,
     project_name: Optional[str] = None,
+    project_id: Optional[str] = None,
     slug: Optional[str] = None,
     global_function: Optional[str] = None,
     # arguments to the function
     input: Any = None,
     messages: Optional[List[Any]] = None,
+    context: Optional[InvokeContext] = None,
     metadata: Optional[Dict[str, Any]] = None,
     tags: Optional[List[str]] = None,
     parent: Optional[Union[Exportable, str]] = None,
@@ -93,6 +116,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,6 +143,8 @@ 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.
 
@@ -161,12 +188,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:

braintrust/generated_types.py

@@ -1,4 +1,4 @@
-"""Auto-generated file (internal git SHA 8e9c0a96b3cf291360978c17580f72f6817bd6c8) -- do not modify"""
+"""Auto-generated file (internal git SHA 437eb5379a737f70dec98033fccf81de43e8e177) -- 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,14 @@ from ._generated_types import (
     GraphNode,
     Group,
     IfExists,
+    InvokeContext,
     InvokeFunction,
     InvokeParent,
+    InvokeScope,
+    MCPServer,
     MessageRole,
     ModelParams,
+    NullableSavedFunctionId,
     ObjectReference,
     ObjectReferenceNullish,
     OnlineScoreConfig,
@@ -86,11 +91,13 @@ from ._generated_types import (
     ServiceToken,
     SpanAttributes,
     SpanIFrame,
+    SpanScope,
     SpanType,
     SSEConsoleEventData,
     SSEProgressEventData,
     StreamingMode,
     ToolFunctionDefinition,
+    TraceScope,
     UploadStatus,
     User,
     View,
@@ -131,6 +138,7 @@ __all__ = [
     "ExperimentEvent",
     "ExtendedSavedFunctionId",
     "ExternalAttachmentReference",
+    "FacetData",
     "Function",
     "FunctionData",
     "FunctionFormat",
@@ -146,10 +154,14 @@ __all__ = [
     "GraphNode",
     "Group",
     "IfExists",
+    "InvokeContext",
     "InvokeFunction",
     "InvokeParent",
+    "InvokeScope",
+    "MCPServer",
     "MessageRole",
     "ModelParams",
+    "NullableSavedFunctionId",
     "ObjectReference",
     "ObjectReferenceNullish",
     "OnlineScoreConfig",
@@ -187,9 +199,11 @@ __all__ = [
     "ServiceToken",
     "SpanAttributes",
     "SpanIFrame",
+    "SpanScope",
     "SpanType",
     "StreamingMode",
     "ToolFunctionDefinition",
+    "TraceScope",
     "UploadStatus",
     "User",
     "View",

braintrust/oai.py

@@ -1,8 +1,10 @@
 import abc
+import base64
+import re
 import time
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Union
 
-from .logger import Span, start_span
+from .logger import Attachment, Span, start_span
 from .span_types import SpanTypeAttribute
 from .util import merge_dicts
 
@@ -68,6 +70,75 @@ def log_headers(response: Any, span: Span):
         )
 
 
+def _convert_data_url_to_attachment(data_url: str, filename: Optional[str] = None) -> Union[Attachment, str]:
+    """Helper function to convert data URL to an Attachment."""
+    data_url_match = re.match(r"^data:([^;]+);base64,(.+)$", data_url)
+    if not data_url_match:
+        return data_url
+
+    mime_type, base64_data = data_url_match.groups()
+
+    try:
+        binary_data = base64.b64decode(base64_data)
+
+        if filename is None:
+            extension = mime_type.split("/")[1] if "/" in mime_type else "bin"
+            prefix = "image" if mime_type.startswith("image/") else "document"
+            filename = f"{prefix}.{extension}"
+
+        attachment = Attachment(data=binary_data, filename=filename, content_type=mime_type)
+
+        return attachment
+    except Exception:
+        return data_url
+
+
+def _process_attachments_in_input(input_data: Any) -> Any:
+    """Process input to convert data URL images and base64 documents to Attachment objects."""
+    if isinstance(input_data, list):
+        return [_process_attachments_in_input(item) for item in input_data]
+
+    if isinstance(input_data, dict):
+        # Check for OpenAI's image_url format with data URLs
+        if (
+            input_data.get("type") == "image_url"
+            and isinstance(input_data.get("image_url"), dict)
+            and isinstance(input_data["image_url"].get("url"), str)
+        ):
+            processed_url = _convert_data_url_to_attachment(input_data["image_url"]["url"])
+            return {
+                **input_data,
+                "image_url": {
+                    **input_data["image_url"],
+                    "url": processed_url,
+                },
+            }
+
+        # Check for OpenAI's file format with data URL (e.g., PDFs)
+        if (
+            input_data.get("type") == "file"
+            and isinstance(input_data.get("file"), dict)
+            and isinstance(input_data["file"].get("file_data"), str)
+        ):
+            file_filename = input_data["file"].get("filename")
+            processed_file_data = _convert_data_url_to_attachment(
+                input_data["file"]["file_data"],
+                filename=file_filename if isinstance(file_filename, str) else None,
+            )
+            return {
+                **input_data,
+                "file": {
+                    **input_data["file"],
+                    "file_data": processed_file_data,
+                },
+            }
+
+        # Recursively process nested objects
+        return {key: _process_attachments_in_input(value) for key, value in input_data.items()}
+
+    return input_data
+
+
 class ChatCompletionWrapper:
     def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
         self.create_fn = create_fn
@@ -190,10 +261,14 @@ class ChatCompletionWrapper:
         # Then, copy the rest of the params
         params = prettify_params(params)
         messages = params.pop("messages", None)
+
+        # Process attachments in input (convert data URLs to Attachment objects)
+        processed_input = _process_attachments_in_input(messages)
+
         return merge_dicts(
             ret,
             {
-                "input": messages,
+                "input": processed_input,
                 "metadata": {**params, "provider": "openai"},
             },
         )
@@ -379,10 +454,14 @@ class ResponseWrapper:
         # Then, copy the rest of the params
         params = prettify_params(params)
         input_data = params.pop("input", None)
+
+        # Process attachments in input (convert data URLs to Attachment objects)
+        processed_input = _process_attachments_in_input(input_data)
+
         return merge_dicts(
             ret,
             {
-                "input": input_data,
+                "input": processed_input,
                 "metadata": {**params, "provider": "openai"},
             },
         )
@@ -540,12 +619,15 @@ class BaseWrapper(abc.ABC):
         ret = params.pop("span_info", {})
 
         params = prettify_params(params)
-        input = params.pop("input", None)
+        input_data = params.pop("input", None)
+
+        # Process attachments in input (convert data URLs to Attachment objects)
+        processed_input = _process_attachments_in_input(input_data)
 
         return merge_dicts(
             ret,
             {
-                "input": input,
+                "input": processed_input,
                 "metadata": {**params, "provider": "openai"},
             },
         )

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.3.14"
+VERSION = "0.3.15"
 
 # this will be templated during the build
-GIT_COMMIT = "dbbc1894ef31143816e5913676301261bc44aa4c"
+GIT_COMMIT = "dcd4f5a4be171b1cac28a5eb3534e4b55420cc06"