braintrust 0.3.13__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/gitutil.py

@@ -88,8 +88,12 @@ def get_past_n_ancestors(n=1000, remote=None):
         if ancestor_output is None:
             return
         ancestor = repo.commit(ancestor_output)
+        count = 0
         for _ in range(n):
+            if count >= n:
+                break
             yield ancestor.hexsha
+            count += 1
             try:
                 if ancestor.parents:
                     ancestor = ancestor.parents[0]

braintrust/logger.py

@@ -1104,7 +1104,7 @@ class _HTTPBackgroundLogger:
             _HTTPBackgroundLogger._write_payload_to_dir(payload_dir=self.all_publish_payloads_dir, payload=dataStr)
         for i in range(self.num_tries):
             start_time = time.time()
-            resp = conn.post("/logs3", data=dataStr)
+            resp = conn.post("/logs3", data=dataStr.encode("utf-8"))
             if resp.ok:
                 return
             resp_errmsg = f"{resp.status_code}: {resp.text}"

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/score.py

@@ -34,6 +34,7 @@ class Score(SerializableDataClass):
 
     def as_dict(self):
         return {
+            "name": self.name,
             "score": self.score,
             "metadata": self.metadata,
         }

braintrust/test_score.py

@@ -0,0 +1,157 @@
+import json
+import unittest
+
+from .score import Score
+
+
+class TestScore(unittest.TestCase):
+    def test_as_dict_includes_all_required_fields(self):
+        """Test that as_dict() includes name, score, and metadata fields."""
+        score = Score(name="test_scorer", score=0.85, metadata={"key": "value"})
+        result = score.as_dict()
+
+        self.assertIn("name", result)
+        self.assertIn("score", result)
+        self.assertIn("metadata", result)
+
+        self.assertEqual(result["name"], "test_scorer")
+        self.assertEqual(result["score"], 0.85)
+        self.assertEqual(result["metadata"], {"key": "value"})
+
+    def test_as_dict_with_null_score(self):
+        """Test that as_dict() works correctly with null score."""
+        score = Score(name="null_scorer", score=None, metadata={})
+        result = score.as_dict()
+
+        self.assertEqual(result["name"], "null_scorer")
+        self.assertIsNone(result["score"])
+        self.assertEqual(result["metadata"], {})
+
+    def test_as_dict_with_empty_metadata(self):
+        """Test that as_dict() works correctly with empty metadata."""
+        score = Score(name="empty_metadata_scorer", score=1.0)
+        result = score.as_dict()
+
+        self.assertEqual(result["name"], "empty_metadata_scorer")
+        self.assertEqual(result["score"], 1.0)
+        self.assertEqual(result["metadata"], {})
+
+    def test_as_dict_with_complex_metadata(self):
+        """Test that as_dict() works correctly with complex nested metadata."""
+        complex_metadata = {
+            "reason": "Test reason",
+            "details": {"nested": {"deeply": "value"}},
+            "list": [1, 2, 3],
+            "bool": True,
+        }
+        score = Score(name="complex_scorer", score=0.5, metadata=complex_metadata)
+        result = score.as_dict()
+
+        self.assertEqual(result["name"], "complex_scorer")
+        self.assertEqual(result["score"], 0.5)
+        self.assertEqual(result["metadata"], complex_metadata)
+
+    def test_as_json_serialization(self):
+        """Test that as_json() produces valid JSON string."""
+        score = Score(name="json_scorer", score=0.75, metadata={"test": "data"})
+        json_str = score.as_json()
+
+        # Should be valid JSON
+        parsed = json.loads(json_str)
+
+        self.assertEqual(parsed["name"], "json_scorer")
+        self.assertEqual(parsed["score"], 0.75)
+        self.assertEqual(parsed["metadata"], {"test": "data"})
+
+    def test_from_dict_round_trip(self):
+        """Test that Score can be serialized to dict and deserialized back."""
+        original = Score(
+            name="round_trip_scorer", score=0.95, metadata={"info": "test"}
+        )
+
+        # Serialize to dict
+        as_dict = original.as_dict()
+
+        # Deserialize from dict
+        restored = Score.from_dict(as_dict)
+
+        self.assertEqual(restored.name, original.name)
+        self.assertEqual(restored.score, original.score)
+        self.assertEqual(restored.metadata, original.metadata)
+
+    def test_array_of_scores_serialization(self):
+        """Test that arrays of Score objects can be serialized correctly."""
+        scores = [
+            Score(name="score_1", score=0.8, metadata={"index": 1}),
+            Score(name="score_2", score=0.6, metadata={"index": 2}),
+            Score(name="score_3", score=None, metadata={}),
+        ]
+
+        # Serialize each score
+        serialized = [s.as_dict() for s in scores]
+
+        # Check that all scores have required fields
+        for i, s_dict in enumerate(serialized):
+            self.assertIn("name", s_dict)
+            self.assertIn("score", s_dict)
+            self.assertIn("metadata", s_dict)
+            self.assertEqual(s_dict["name"], f"score_{i + 1}")
+
+        # Check specific values
+        self.assertEqual(serialized[0]["score"], 0.8)
+        self.assertEqual(serialized[1]["score"], 0.6)
+        self.assertIsNone(serialized[2]["score"])
+
+    def test_array_of_scores_json_serialization(self):
+        """Test that arrays of Score objects can be JSON serialized."""
+        scores = [
+            Score(name="json_score_1", score=0.9),
+            Score(name="json_score_2", score=0.7),
+        ]
+
+        # Serialize to JSON
+        serialized = [s.as_dict() for s in scores]
+        json_str = json.dumps(serialized)
+
+        # Parse back
+        parsed = json.loads(json_str)
+
+        self.assertEqual(len(parsed), 2)
+        self.assertEqual(parsed[0]["name"], "json_score_1")
+        self.assertEqual(parsed[0]["score"], 0.9)
+        self.assertEqual(parsed[1]["name"], "json_score_2")
+        self.assertEqual(parsed[1]["score"], 0.7)
+
+    def test_score_validation_enforces_bounds(self):
+        """Test that Score validates score values are between 0 and 1."""
+        # Valid scores
+        Score(name="valid_0", score=0.0)
+        Score(name="valid_1", score=1.0)
+        Score(name="valid_mid", score=0.5)
+        Score(name="valid_null", score=None)
+
+        # Invalid scores
+        with self.assertRaises(ValueError):
+            Score(name="invalid_negative", score=-0.1)
+
+        with self.assertRaises(ValueError):
+            Score(name="invalid_over_one", score=1.1)
+
+    def test_score_does_not_include_deprecated_error_field(self):
+        """Test that as_dict() does not include the deprecated error field."""
+        score = Score(name="test_scorer", score=0.5)
+        result = score.as_dict()
+
+        # The error field should not be in the serialized output
+        self.assertNotIn("error", result)
+
+        # Even if error was set (though deprecated), it shouldn't be in as_dict
+        score_with_error = Score(name="error_scorer", score=0.5)
+        score_with_error.error = Exception("test")  # Set after construction
+        result_with_error = score_with_error.as_dict()
+
+        self.assertNotIn("error", result_with_error)
+
+
+if __name__ == "__main__":
+    unittest.main()

braintrust/version.py

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