braintrust 0.5.3__py3-none-any.whl → 0.5.4__py3-none-any.whl

braintrust/_generated_types.py

@@ -194,6 +194,10 @@ class BatchedFacetDataFacet(TypedDict):
     """
     The model to use for facet extraction
     """
+    embedding_model: NotRequired[str | None]
+    """
+    The embedding model to use for vectorizing facet results.
+    """
     no_match_pattern: NotRequired[str | None]
     """
     Regex pattern to identify outputs that do not match the facet. If the output matches, the facet will be saved as 'no_match'
@@ -740,9 +744,32 @@ class FunctionDataFunctionData2(TypedDict):
     endpoint: str
     eval_name: str
     parameters: Mapping[str, Any]
+    parameters_version: NotRequired[str | None]
+    """
+    The version (transaction ID) of the parameters being used
+    """
+
+
+class FunctionDataFunctionData4Schema(TypedDict):
+    type: Literal['object']
+    properties: Mapping[str, Mapping[str, Any]]
+    required: NotRequired[Sequence[str] | None]
+    additionalProperties: NotRequired[bool | None]
+
+
+class FunctionDataFunctionData4(TypedDict):
+    type: Literal['parameters']
+    data: Mapping[str, Any]
+    """
+    The parameters data
+    """
+    __schema: FunctionDataFunctionData4Schema
+    """
+    JSON Schema format for parameters
+    """
 
 
-FunctionFormat: TypeAlias = Literal['llm', 'code', 'global', 'graph']
+FunctionFormat: TypeAlias = Literal['llm', 'code', 'global', 'graph', 'topic_map']
 
 
 class FunctionIdFunctionId(TypedDict):
@@ -807,7 +834,7 @@ FunctionIdRef: TypeAlias = Mapping[str, Any]
 
 
 FunctionObjectType: TypeAlias = Literal[
-    'prompt', 'tool', 'scorer', 'task', 'workflow', 'custom_view', 'preprocessor', 'facet', 'classifier'
+    'prompt', 'tool', 'scorer', 'task', 'workflow', 'custom_view', 'preprocessor', 'facet', 'classifier', 'parameters'
 ]
 
 
@@ -815,7 +842,7 @@ FunctionOutputType: TypeAlias = Literal['completion', 'score', 'facet', 'classif
 
 
 FunctionTypeEnum: TypeAlias = Literal[
-    'llm', 'scorer', 'task', 'tool', 'custom_view', 'preprocessor', 'facet', 'classifier', 'tag'
+    'llm', 'scorer', 'task', 'tool', 'custom_view', 'preprocessor', 'facet', 'classifier', 'tag', 'parameters'
 ]
 """
 The type of global function. Defaults to 'scorer'.
@@ -823,7 +850,7 @@ The type of global function. Defaults to 'scorer'.
 
 
 FunctionTypeEnumNullish: TypeAlias = Literal[
-    'llm', 'scorer', 'task', 'tool', 'custom_view', 'preprocessor', 'facet', 'classifier', 'tag'
+    'llm', 'scorer', 'task', 'tool', 'custom_view', 'preprocessor', 'facet', 'classifier', 'tag', 'parameters'
 ]
 
 
@@ -2236,6 +2263,64 @@ class ToolFunctionDefinition(TypedDict):
     function: ToolFunctionDefinitionFunction
 
 
+class TopicMapReportSettings(TypedDict):
+    algorithm: Literal['hdbscan', 'kmeans', 'hierarchical']
+    dimension_reduction: Literal['umap', 'pca', 'none']
+    vector_field: str
+    embedding_model: str
+    n_clusters: NotRequired[int | None]
+    umap_dimensions: NotRequired[int | None]
+    min_cluster_size: NotRequired[int | None]
+    min_samples: NotRequired[int | None]
+
+
+class TopicMapReportQuerySettings(TypedDict):
+    hierarchy_threshold: NotRequired[int | None]
+    auto_naming: NotRequired[bool | None]
+    skip_cache: NotRequired[bool | None]
+    viz_mode: NotRequired[Literal['bar', 'scatter'] | None]
+    naming_model: NotRequired[str | None]
+
+
+class TopicMapReportClusterSample(TypedDict):
+    id: str
+    text: str
+    root_span_id: str
+    span_id: str
+
+
+class TopicMapReportCluster(TypedDict):
+    cluster_id: float
+    parent_cluster_id: NotRequired[float | None]
+    topic_id: str
+    count: float
+    sample_texts: Sequence[str]
+    samples: Sequence[TopicMapReportClusterSample]
+    name: NotRequired[str | None]
+    description: NotRequired[str | None]
+    keywords: NotRequired[Sequence[str] | None]
+    centroid: NotRequired[Sequence[float] | None]
+    parent_id: NotRequired[float | None]
+    is_leaf: NotRequired[bool | None]
+    depth: NotRequired[float | None]
+
+
+class TopicMapReportEmbeddingPoint(TypedDict):
+    x: float
+    y: float
+    cluster: float
+    text: NotRequired[str | None]
+
+
+class TopicMapReport(TypedDict):
+    version: Literal[1]
+    created_at: NotRequired[str | None]
+    settings: TopicMapReportSettings
+    query_settings: TopicMapReportQuerySettings
+    clusters: Sequence[TopicMapReportCluster]
+    embedding_points: NotRequired[Sequence[TopicMapReportEmbeddingPoint] | None]
+
+
 class TraceScope(TypedDict):
     type: Literal['trace']
     idle_seconds: NotRequired[float | None]
@@ -2500,12 +2585,6 @@ class PreprocessorPreprocessor4(PreprocessorPreprocessor1, PreprocessorPreproces
 Preprocessor: TypeAlias = PreprocessorPreprocessor3 | PreprocessorPreprocessor4
 
 
-class BatchedFacetData(TypedDict):
-    type: Literal['batched_facet']
-    preprocessor: NotRequired[Preprocessor | None]
-    facets: Sequence[BatchedFacetDataFacet]
-
-
 ChatCompletionContentPart: TypeAlias = (
     ChatCompletionContentPartTextWithTitle
     | ChatCompletionContentPartImageWithTitle
@@ -2761,6 +2840,10 @@ class FacetData(TypedDict):
     """
     The model to use for facet extraction
     """
+    embedding_model: NotRequired[str | None]
+    """
+    The embedding model to use for vectorizing facet results.
+    """
     no_match_pattern: NotRequired[str | None]
     """
     Regex pattern to identify outputs that do not match the facet. If the output matches, the facet will be saved as 'no_match'
@@ -3044,11 +3127,54 @@ class SpanAttributes(TypedDict):
     type: NotRequired[SpanType | None]
 
 
+class TopicMapData(TypedDict):
+    type: Literal['topic_map']
+    source_facet: str
+    """
+    The facet field name to use as input for classification
+    """
+    embedding_model: str
+    """
+    The embedding model to use for embedding facet values
+    """
+    bundle_key: str
+    """
+    Key of the topic map bundle in code_bundles bucket
+    """
+    distance_threshold: NotRequired[float | None]
+    """
+    Maximum distance to nearest centroid. If exceeded, returns no_match.
+    """
+    report: NotRequired[TopicMapReport | None]
+
+
 class ViewData(TypedDict):
     search: NotRequired[ViewDataSearch | None]
     custom_charts: NotRequired[Any | None]
 
 
+class BatchedFacetDataTopicMaps(TypedDict):
+    function_name: str
+    """
+    The name of the topic map function
+    """
+    topic_map_id: NotRequired[str | None]
+    """
+    The id of the topic map function
+    """
+    topic_map_data: TopicMapData
+
+
+class BatchedFacetData(TypedDict):
+    type: Literal['batched_facet']
+    preprocessor: NotRequired[Preprocessor | None]
+    facets: Sequence[BatchedFacetDataFacet]
+    topic_maps: NotRequired[Mapping[str, BatchedFacetDataTopicMaps] | None]
+    """
+    Topic maps that depend on facets in this batch, keyed by source facet name
+    """
+
+
 class ExperimentEvent(TypedDict):
     id: str
     """
@@ -3376,6 +3502,7 @@ class View(TypedDict):
         'datasets',
         'dataset',
         'prompts',
+        'parameters',
         'tools',
         'scorers',
         'classifiers',
@@ -3617,6 +3744,8 @@ FunctionData: TypeAlias = (
     | FunctionDataFunctionData3
     | FacetData
     | BatchedFacetData
+    | FunctionDataFunctionData4
+    | TopicMapData
 )
 
 

braintrust/generated_types.py

@@ -1,4 +1,4 @@
-"""Auto-generated file (internal git SHA 7178200dd3c7869f27b677260303cda0b798bf42) -- do not modify"""
+"""Auto-generated file (internal git SHA 2822255bed426d5442adc880a8f71b8a378de3d4) -- do not modify"""
 
 from ._generated_types import (
     Acl,
@@ -101,6 +101,8 @@ from ._generated_types import (
     SSEProgressEventData,
     StreamingMode,
     ToolFunctionDefinition,
+    TopicMapData,
+    TopicMapReport,
     TraceScope,
     TriggeredFunctionState,
     UploadStatus,
@@ -212,6 +214,8 @@ __all__ = [
     "SpanType",
     "StreamingMode",
     "ToolFunctionDefinition",
+    "TopicMapData",
+    "TopicMapReport",
     "TraceScope",
     "TriggeredFunctionState",
     "UploadStatus",

braintrust/oai.py

@@ -312,16 +312,25 @@ class ChatCompletionWrapper:
 
                 # pylint: disable=unsubscriptable-object
                 if not tool_calls or (tool_delta.get("id") and tool_calls[-1]["id"] != tool_delta.get("id")):
+                    function_arg = tool_delta.get("function", {})
                     tool_calls = (tool_calls or []) + [
                         {
-                            "id": delta["tool_calls"][0]["id"],
-                            "type": delta["tool_calls"][0]["type"],
-                            "function": delta["tool_calls"][0]["function"],
+                            "id": tool_delta.get("id"),
+                            "type": tool_delta.get("type"),
+                            "function": {
+                                "name": function_arg.get("name"),
+                                "arguments": function_arg.get("arguments") or "",
+                            },
                         }
                     ]
                 else:
                     # pylint: disable=unsubscriptable-object
-                    tool_calls[-1]["function"]["arguments"] += delta["tool_calls"][0]["function"]["arguments"]
+                    # append to existing tool call
+                    function_arg = tool_delta.get("function", {})
+                    args = function_arg.get("arguments") or ""
+                    if isinstance(args, str):
+                        # pylint: disable=unsubscriptable-object
+                        tool_calls[-1]["function"]["arguments"] += args
 
         return {
             "metrics": metrics,

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.5.3"
+VERSION = "0.5.4"
 
 # this will be templated during the build
-GIT_COMMIT = "ce261c13723e8985e61819f6df56298f863155d0"
+GIT_COMMIT = "449f6eeddd54f73d88e185fcf499087d3e61fbe6"

braintrust/wrappers/test_openai.py

@@ -5,6 +5,7 @@ import braintrust
 import openai
 import pytest
 from braintrust import logger, wrap_openai
+from braintrust.oai import ChatCompletionWrapper
 from braintrust.test_helpers import assert_dict_matches, init_test_logger
 from braintrust.wrappers.test_utils import assert_metrics_are_valid, run_in_subprocess, verify_autoinstrument_script
 from openai import AsyncOpenAI
@@ -377,7 +378,6 @@ def test_openai_responses_sparse_indices(memory_logger):
     # No spans should be generated from this unit test
     assert not memory_logger.pop()
 
-
 @pytest.mark.vcr
 def test_openai_embeddings(memory_logger):
     assert not memory_logger.pop()
@@ -1933,3 +1933,102 @@ class TestAutoInstrumentOpenAI:
     def test_auto_instrument_openai(self):
         """Test auto_instrument patches OpenAI, creates spans, and uninstrument works."""
         verify_autoinstrument_script("test_auto_openai.py")
+
+class TestZAICompatibleOpenAI:
+    """Tests for validating some ZAI compatibility with OpenAI wrapper."""
+
+    def test_chat_completion_streaming_none_arguments(self, memory_logger):
+        """Test that ChatCompletionWrapper handles None arguments in tool calls (e.g., GLM-4.6 behavior)."""
+        assert not memory_logger.pop()
+
+        # Simulate streaming results with None arguments in tool calls
+        # This mimics the behavior of GLM-4.6 which returns {'arguments': None, 'name': 'weather'}
+        all_results = [
+            # First chunk: initial tool call with None arguments
+            {
+                "choices": [
+                    {
+                        "delta": {
+                            "role": "assistant",
+                            "tool_calls": [
+                                {
+                                    "id": "call_123",
+                                    "type": "function",
+                                    "function": {
+                                        "name": "get_weather",
+                                        "arguments": None,  # GLM-4.6 returns None here
+                                    },
+                                }
+                            ],
+                        },
+                        "finish_reason": None,
+                    }
+                ],
+            },
+            # Second chunk: subsequent tool call arguments (also None)
+            {
+                "choices": [
+                    {
+                        "delta": {
+                            "tool_calls": [
+                                {
+                                    "function": {
+                                        "arguments": None,  # Subsequent chunks can also have None
+                                    }
+                                }
+                            ],
+                        },
+                        "finish_reason": None,
+                    }
+                ],
+            },
+            # Third chunk: actual arguments
+            {
+                "choices": [
+                    {
+                        "delta": {
+                            "tool_calls": [
+                                {
+                                    "function": {
+                                        "arguments": '{"city": "New York"}',
+                                    }
+                                }
+                            ],
+                        },
+                        "finish_reason": None,
+                    }
+                ],
+            },
+            # Final chunk
+            {
+                "choices": [
+                    {
+                        "delta": {},
+                        "finish_reason": "tool_calls",
+                    }
+                ],
+            },
+        ]
+
+        # Process the results
+        wrapper = ChatCompletionWrapper(None, None)
+        result = wrapper._postprocess_streaming_results(all_results)
+
+        # Verify the output was built correctly
+        assert "output" in result
+        assert len(result["output"]) == 1
+        message = result["output"][0]["message"]
+        assert message["role"] == "assistant"
+        assert message["tool_calls"] is not None
+        assert len(message["tool_calls"]) == 1
+
+        # Verify the tool call was assembled correctly despite None arguments
+        tool_call = message["tool_calls"][0]
+        assert tool_call["id"] == "call_123"
+        assert tool_call["type"] == "function"
+        assert tool_call["function"]["name"] == "get_weather"
+        # The arguments should be the concatenation: "" + "" + '{"city": "New York"}'
+        assert tool_call["function"]["arguments"] == '{"city": "New York"}'
+
+        # No spans should be generated from this unit test
+        assert not memory_logger.pop()

{braintrust-0.5.3.dist-info → braintrust-0.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrust
-Version: 0.5.3
+Version: 0.5.4
 Summary: SDK for integrating Braintrust
 Home-page: https://www.braintrust.dev
 Author: Braintrust

{braintrust-0.5.3.dist-info → braintrust-0.5.4.dist-info}/RECORD

@@ -1,5 +1,5 @@
 braintrust/__init__.py,sha256=DU4gzkV0R3nsWxp2da7iQS1MM_T9xHkrKSQE6nAnZbE,2627
-braintrust/_generated_types.py,sha256=1gNbTwsy7LEBdLW-uYqzQ2v78ToAFX0PSmeE_NOJC18,101153
+braintrust/_generated_types.py,sha256=d1z6xRHPc0iqifiaoCh7YwDQNSxKZAhM2KYmG0Fyrvk,104819
 braintrust/audit.py,sha256=3GQKzuTcFquYdrJtABM-k3xMlOIqgVkfG6UyeQ8_028,461
 braintrust/auto.py,sha256=wf4Jb7hYoGS0Agpx-YjUYEW7wwwUpyAJDp-l3A-y6c0,4792
 braintrust/aws.py,sha256=OBz_SRyopgpCDSNvETLypzGwTXk-bNLn-Eisevnjfwo,377
@@ -9,14 +9,14 @@ braintrust/context.py,sha256=bOo1Li29lUsi2DOOllIDar0oRuQNbkLNzJ7cYq5JTbo,4126
 braintrust/db_fields.py,sha256=AX-K5t7KqO-xHHOfVRv8bn1ww7gZd3RNIFfANkZ2W0U,709
 braintrust/framework.py,sha256=KoXFKCprfeEHq-AqjarENcmHHcTuvr4PScWtMbIQ6zg,62635
 braintrust/framework2.py,sha256=o0igz4vXbmn0jHJPhDYvx14rFnI3ntV8H6VJfyJYRtM,16542
-braintrust/generated_types.py,sha256=MUpi5l9XIW0bGqnYqa9DE7XbxDYsBI82ANyTVHz9DVg,5047
+braintrust/generated_types.py,sha256=WExtraVBQZMS5DcCVSQl_k66p6VUY7cK7u1SSQhc_to,5127
 braintrust/git_fields.py,sha256=au5ayyuvt7y_ojE9LC98ypTZd3RgFdjhRc8eFxcjnto,1434
 braintrust/gitutil.py,sha256=RsW7cawJMgaAaTw6WeB1sShyfflkPb7yH2x7yuRv10c,5642
 braintrust/http_headers.py,sha256=9ZsDcsAKG04SGowsgchZktD6rG_oSTKWa8QyGUPA4xE,154
 braintrust/id_gen.py,sha256=4UWLWRhksf76IkYi4cKACSaQ3yNgausrMRlhiurhy74,1590
 braintrust/logger.py,sha256=ofx0_8ywVk_iyyqOH-6kCgRaYmm_sO7guVJ5hW2ISNg,217467
 braintrust/merge_row_batch.py,sha256=mCutumLDOpH8ArP_4K5swP93mve7gmgfAnQdkdjUZ5E,6271
-braintrust/oai.py,sha256=K7NLw7-3U7TJCyBVeoOwpy_UOmGlpaXXJcHZgW-rzHM,37746
+braintrust/oai.py,sha256=ispeHAeykwF3VmGMnaVKZRn7Q80DMtpM7nitFrtybRQ,38193
 braintrust/object.py,sha256=vYLyYWncsqLD00zffZUJwGTSkcJF9IIXmgIzrx3Np5c,632
 braintrust/parameters.py,sha256=sQWfw18QXdPSnMHsF7aRrPmP7Zx6HEz9vaTUXWreudg,5911
 braintrust/prompt.py,sha256=pLzhXoBp7ebxNraZADflR4YQMW5Ycjmt5ucK80P4_h0,1875
@@ -51,7 +51,7 @@ braintrust/test_util.py,sha256=SuSKTmvNyaR9Rbgf2TYCUWxJpZHoA2BMx6n4nQfV_pM,8221
 braintrust/test_version.py,sha256=hk5JKjEFbNJ_ONc1VEkqHquflzre34RpFhCEYLTK8iA,1051
 braintrust/trace.py,sha256=PHxfaHApGP_MPMIndZw7atIa2iwKPETUoG-TbF2dv6A,13767
 braintrust/util.py,sha256=S-qMBNsT36r_3pJ4LNKZ-vvHRlJwy8Wy7M7NAdfNOug,8919
-braintrust/version.py,sha256=u7mu2r_lw0F22zAYs8OHbW11HrkdIhoAK-RbDs2uQdA,117
+braintrust/version.py,sha256=9Ag0YVXSTg_ksT5pb04Gdog87NeKSiNXtmslh5ZQyhQ,117
 braintrust/xact_ids.py,sha256=bdyp88HjlyIkglgLSqYlCYscdSH6EWVyE14sR90Xl1s,658
 braintrust/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 braintrust/cli/__main__.py,sha256=wCBKHGVmn3IT_yMXk5qfDwyI2SV2gf1tLr0NTxm9T8k,1519
@@ -109,7 +109,7 @@ braintrust/wrappers/test_google_genai.py,sha256=BKkPfLFk2JG1UNJ3mFKzv5OjioWqo7EQ
 braintrust/wrappers/test_langsmith_wrapper.py,sha256=wEbPNy4o7VVvcuHcsCJ-sy2EATvBxhUXTYFBQNkKCjs,10449
 braintrust/wrappers/test_litellm.py,sha256=MKukVH-C-Mwc0-cVJZhkLdXwiXupXHc8EiWYKq0X8V0,23620
 braintrust/wrappers/test_oai_attachments.py,sha256=_EtNXjQxPgqXmj6UYMZn9GF4GDZf8m_1_TrwiEk7HWQ,11100
-braintrust/wrappers/test_openai.py,sha256=pC-5zO2CbeG1wmQ5IvbCQ5vXic5sZgTCTW7347QwbwI,69474
+braintrust/wrappers/test_openai.py,sha256=dBphKJT2-xSnK4mJ_xbKLFPBfpHzK07TmHxPTQ8ZQDg,73355
 braintrust/wrappers/test_openrouter.py,sha256=8HUfILPugOMqcvttpq76KQrynFb0xZpazvta7TTSF6A,3849
 braintrust/wrappers/test_pydantic_ai_integration.py,sha256=xUSkiY5HUI-Z9G_etjyX7rGlOjBEdEY9CwVHyWM3xEE,104460
 braintrust/wrappers/test_pydantic_ai_wrap_openai.py,sha256=OO5NrbothkMr4v2sZ-EZLH7-yLj3k6TfdLG4uzXAsQk,5090
@@ -125,8 +125,8 @@ braintrust/wrappers/claude_agent_sdk/__init__.py,sha256=4FUE59ii39jVfhMAfkOcU-TJ
 braintrust/wrappers/claude_agent_sdk/_wrapper.py,sha256=UNZBcgTu7X71zvJKy6e-QQFuz9j8kRS6Kd1VOh3OOXA,17755
 braintrust/wrappers/claude_agent_sdk/test_wrapper.py,sha256=RmSzTfDC3tkepEbc9S_KYFITsjGVyllYE1fF9roSqBk,10779
 braintrust/wrappers/google_genai/__init__.py,sha256=C7MzKUr17CEcIbP0kg9L9IjN-6suE3NQ-oSYAjXdR_g,15832
-braintrust-0.5.3.dist-info/METADATA,sha256=snLiF8lw3wJkTkmgZnbthtfApgwnFb7eKFNARCWK0mM,3753
-braintrust-0.5.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-braintrust-0.5.3.dist-info/entry_points.txt,sha256=Zpc0_09g5xm8as5jHqqFq7fhwO0xHSNct_TrEMONS7Q,60
-braintrust-0.5.3.dist-info/top_level.txt,sha256=hw1-y-UFMf60RzAr8x_eM7SThbIuWfQsQIbVvqSF83A,11
-braintrust-0.5.3.dist-info/RECORD,,
+braintrust-0.5.4.dist-info/METADATA,sha256=Z3DmXXIkkIpZ5XQ7I6ECeVg5G3SW29onbNWyoLlvOZM,3753
+braintrust-0.5.4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+braintrust-0.5.4.dist-info/entry_points.txt,sha256=Zpc0_09g5xm8as5jHqqFq7fhwO0xHSNct_TrEMONS7Q,60
+braintrust-0.5.4.dist-info/top_level.txt,sha256=hw1-y-UFMf60RzAr8x_eM7SThbIuWfQsQIbVvqSF83A,11
+braintrust-0.5.4.dist-info/RECORD,,