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

braintrust/test_merge_row_batch.py

@@ -0,0 +1,160 @@
+import unittest
+
+from braintrust.db_fields import IS_MERGE_FIELD
+from braintrust.merge_row_batch import batch_items, merge_row_batch
+
+
+class MergeRowBatchTest(unittest.TestCase):
+    def test_basic(self):
+        rows = [
+            # These rows should get merged together, ending up as a merge.
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(a=12),
+                **{IS_MERGE_FIELD: True},
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(b=10),
+                **{IS_MERGE_FIELD: True},
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(c="hello"),
+                **{IS_MERGE_FIELD: True},
+            ),
+            # The first row should be clobbered by the second, but the third
+            # merged with the second, ending up as a replacement.
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(a="hello"),
+            ),
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(b=10),
+            ),
+            dict(
+                experiment_id="e0",
+                id="y",
+                inputs=dict(c=12),
+                **{IS_MERGE_FIELD: True},
+            ),
+            # These rows should be clobbered separately from the last batch.
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(a="hello"),
+            ),
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(b=10),
+            ),
+            dict(
+                dataset_id="d0",
+                id="y",
+                inputs=dict(c=12),
+            ),
+        ]
+
+        merged_rows = merge_row_batch(rows)
+        key_to_rows = {(row.get("experiment_id"), row.get("dataset_id"), row.get("id")): row for row in merged_rows}
+        self.assertEqual(
+            {
+                ("e0", None, "x"): dict(
+                    experiment_id="e0",
+                    id="x",
+                    inputs=dict(a=12, b=10, c="hello"),
+                    **{IS_MERGE_FIELD: True},
+                ),
+                ("e0", None, "y"): dict(
+                    experiment_id="e0",
+                    id="y",
+                    inputs=dict(b=10, c=12),
+                ),
+                (None, "d0", "y"): dict(
+                    dataset_id="d0",
+                    id="y",
+                    inputs=dict(c=12),
+                ),
+            },
+            key_to_rows,
+        )
+
+    def test_skip_fields(self):
+        rows = [
+            # These rows should get merged together, ending up as a merge. But
+            # the original fields should be retained, regardless of whether we
+            # populated them or not.
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(a=12),
+                **{IS_MERGE_FIELD: True},
+                created=123,
+                root_span_id="abc",
+                _parent_id="baz",
+                span_parents=["foo", "bar"],
+            ),
+            dict(
+                experiment_id="e0",
+                id="x",
+                inputs=dict(b=10),
+                **{IS_MERGE_FIELD: True},
+                created=456,
+                span_id="foo",
+                root_span_id="bar",
+                _parent_id="boop",
+                span_parents=[],
+            ),
+        ]
+
+        merged_rows = merge_row_batch(rows)
+        self.assertEqual(
+            merged_rows,
+            [
+                dict(
+                    experiment_id="e0",
+                    id="x",
+                    inputs=dict(a=12, b=10),
+                    **{IS_MERGE_FIELD: True},
+                    created=123,
+                    root_span_id="abc",
+                    _parent_id="baz",
+                    span_parents=["foo", "bar"],
+                ),
+            ],
+        )
+
+
+class BatchItemsTest(unittest.TestCase):
+    def test_basic(self):
+        a = "x" * 1
+        b = "x" * 2
+        c = "x" * 4
+        d = "y" * 1
+        e = "y" * 2
+        f = "y" * 4
+
+        items = [a, b, c, f, e, d]
+
+        # No limits.
+        output = batch_items(items)
+        self.assertEqual(output, [[a, b, c, f, e, d]])
+
+        # Num items limit.
+        output = batch_items(items, batch_max_num_items=2)
+        self.assertEqual(output, [[a, b], [c, f], [e, d]])
+
+        # Num bytes limit.
+        output = batch_items(items, batch_max_num_bytes=2)
+        self.assertEqual(output, [[a], [b], [c], [f], [e], [d]])
+
+        # Both items and num bytes limit.
+        output = batch_items(items, batch_max_num_items=2, batch_max_num_bytes=5)
+        self.assertEqual(output, [[a, b], [c], [f], [e, d]])

braintrust/version.py

@@ -1,4 +1,4 @@
-VERSION = "0.5.2"
+VERSION = "0.5.4"
 
 # this will be templated during the build
-GIT_COMMIT = "25868bc58450dad2058b6499ce3bb9400330fbd1"
+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/wrappers/threads.py

@@ -0,0 +1,114 @@
+import contextvars
+import functools
+import logging
+import threading
+from concurrent import futures
+from typing import Any, TypeVar
+
+from wrapt import wrap_function_wrapper  # pyright: ignore[reportUnknownVariableType, reportMissingTypeStubs]
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["setup_threads", "patch_thread", "patch_thread_pool_executor"]
+
+
+def setup_threads() -> bool:
+    """
+    Setup automatic context propagation for threading.
+
+    This patches stdlib threading primitives to automatically
+    propagate Braintrust context across thread boundaries.
+
+    Enable via:
+        - BRAINTRUST_INSTRUMENT_THREADS=true env var (automatic)
+        - Call this function directly (manual)
+
+    Returns:
+        bool: True if instrumentation was successful, False otherwise.
+    """
+    try:
+        patch_thread(threading.Thread)
+        patch_thread_pool_executor(futures.ThreadPoolExecutor)
+
+        logger.debug("Braintrust thread instrumentation enabled")
+        return True
+
+    except Exception as e:
+        logger.warning(f"Failed to enable thread instrumentation: {e}")
+        return False
+
+
+T = TypeVar("T", bound=type[threading.Thread])
+
+
+def patch_thread(thread_cls: T) -> T:
+    if __is_patched(thread_cls):
+        return thread_cls
+
+    def _wrap_thread_start(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+        try:
+            instance._braintrust_context = contextvars.copy_context()
+        except Exception as e:
+            logger.debug(f"Failed to capture context in thread start: {e}")
+        return wrapped(*args, **kwargs)
+
+    wrap_function_wrapper(thread_cls, "start", _wrap_thread_start)
+
+    def _wrap_thread_run(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+        try:
+            if hasattr(instance, "_braintrust_context"):
+                return instance._braintrust_context.run(wrapped, *args, **kwargs)
+        except Exception as e:
+            logger.debug(f"Failed to restore context in thread run: {e}")
+        return wrapped(*args, **kwargs)
+
+    wrap_function_wrapper(thread_cls, "run", _wrap_thread_run)
+
+    __mark_patched(thread_cls)
+    return thread_cls
+
+
+def __is_patched(obj: Any) -> bool:
+    """Check if an object has already been patched."""
+    return getattr(obj, "_braintrust_patched", False)
+
+
+def __mark_patched(obj: Any) -> None:
+    setattr(obj, "_braintrust_patched", True)
+
+
+P = TypeVar("P", bound=type[futures.ThreadPoolExecutor])
+
+
+def patch_thread_pool_executor(executor_cls: P) -> P:
+    if __is_patched(executor_cls):
+        return executor_cls
+
+    def _wrap_executor_submit(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+        try:
+            if not args:
+                return wrapped(*args, **kwargs)
+
+            func = args[0]
+            ctx = contextvars.copy_context()
+
+            @functools.wraps(func)
+            def context_wrapper(*func_args: Any, **func_kwargs: Any) -> Any:
+                try:
+                    return ctx.run(func, *func_args, **func_kwargs)
+                except Exception as e:
+                    # context.run() can fail if token is invalid
+                    logger.debug(f"Failed to run in captured context: {e}")
+                    return func(*func_args, **func_kwargs)
+
+            new_args = (context_wrapper,) + args[1:]
+            return wrapped(*new_args, **kwargs)
+        except Exception as e:
+            # Wrapping can fail - fall back to original
+            logger.debug(f"Failed to wrap executor submit: {e}")
+            return wrapped(*args, **kwargs)
+
+    wrap_function_wrapper(executor_cls, "submit", _wrap_executor_submit)
+
+    __mark_patched(executor_cls)
+    return executor_cls

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

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

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

@@ -1,23 +1,22 @@
-braintrust/__init__.py,sha256=Bn-JzKyigMCQs9XqjtWQwLoUtConMO-icUCPeIKuYCk,2330
-braintrust/_generated_types.py,sha256=_eqozb1MG75khlsvXoAV7oKjbgEEBvnzlLtiUfpxX8U,99447
+braintrust/__init__.py,sha256=DU4gzkV0R3nsWxp2da7iQS1MM_T9xHkrKSQE6nAnZbE,2627
+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
 braintrust/bt_json.py,sha256=VNunedFUEfbvVEBuACHLRsjDr8lLD3_nFwB54MXtcbY,9385
 braintrust/conftest.py,sha256=qcS1T4tqXN5Ocaol6gEN9xGOXBqvTYPXXiBJjHnHC6Y,3140
 braintrust/context.py,sha256=bOo1Li29lUsi2DOOllIDar0oRuQNbkLNzJ7cYq5JTbo,4126
-braintrust/db_fields.py,sha256=vjVEyDxl2d13lvvTcVwT5paHvaT-1kgI-4Lg-gihRIw,476
+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=_IBZSkYTralTQM5hLhnbdIQS8_NhWr4V77kJPHtEUHY,4997
+braintrust/generated_types.py,sha256=WExtraVBQZMS5DcCVSQl_k66p6VUY7cK7u1SSQhc_to,5127
 braintrust/git_fields.py,sha256=au5ayyuvt7y_ojE9LC98ypTZd3RgFdjhRc8eFxcjnto,1434
 braintrust/gitutil.py,sha256=RsW7cawJMgaAaTw6WeB1sShyfflkPb7yH2x7yuRv10c,5642
-braintrust/graph_util.py,sha256=Z2Uy8RaOq5iMe5mShhQqRDDIpXVitE-biVxDiFB-0Ds,5545
 braintrust/http_headers.py,sha256=9ZsDcsAKG04SGowsgchZktD6rG_oSTKWa8QyGUPA4xE,154
 braintrust/id_gen.py,sha256=4UWLWRhksf76IkYi4cKACSaQ3yNgausrMRlhiurhy74,1590
-braintrust/logger.py,sha256=lxR-4L6HK3PXa_ZQaVTd1-IEzxJ55NE1XmmPFLiwGFw,209535
-braintrust/merge_row_batch.py,sha256=tvuz3qdsa7HZBrzzAzoQqAXU7pKsHUebC7sw6RybbmA,9881
-braintrust/oai.py,sha256=K7NLw7-3U7TJCyBVeoOwpy_UOmGlpaXXJcHZgW-rzHM,37746
+braintrust/logger.py,sha256=ofx0_8ywVk_iyyqOH-6kCgRaYmm_sO7guVJ5hW2ISNg,217467
+braintrust/merge_row_batch.py,sha256=mCutumLDOpH8ArP_4K5swP93mve7gmgfAnQdkdjUZ5E,6271
+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
@@ -33,12 +32,14 @@ braintrust/span_identifier_v3.py,sha256=RAvOK0lK0huH952kI5X1Q9TaAloD5to8jgTCuYMM
 braintrust/span_identifier_v4.py,sha256=uFT-OdzySo4uCeAaJC3VqH55Oy433xZGBdK2hiEsm2w,10044
 braintrust/span_types.py,sha256=cpTzCwUj4yBPbPLnzR23-VXIU2E9nl_dsVCSVMvtSkc,376
 braintrust/test_bt_json.py,sha256=pokqmFSQ3m8aB0XN1EEsS-zbv_N4KCYzad6VyOFtwPw,25188
+braintrust/test_context.py,sha256=4wZOhwzGgAxt-CcN2izAxhgcY975oh_HUMkZXRwwTys,42754
 braintrust/test_framework.py,sha256=rPjJYfVzRSbjf8e0irVfCnE3q5Rl1O_2Dy2dPyC9-lo,18674
 braintrust/test_framework2.py,sha256=pSEEmBIyszAiYnpEVvDZgJqIe3lQ3T807LmIuBqV98w,7235
 braintrust/test_helpers.py,sha256=VSelzjkR2IHyy5QD6sYB_79VIXq-wEDslDwyx9H11kI,13528
 braintrust/test_http.py,sha256=RiW07psmoxO9ZB5ZUN1K1hIan7l4_F9YC2mW0iqSMr0,15751
 braintrust/test_id_gen.py,sha256=Oqp8Rxd1_9ptpEJSu_xi6STIquZpx9n7Kh-bGUKMQH0,2489
 braintrust/test_logger.py,sha256=zINMwUn2hT1_ElZRkVJP95CkTwtLhjUKpR5zIP7tn5Y,112653
+braintrust/test_merge_row_batch.py,sha256=cQ_Hiq6wlWIs091Tqexgd3TXWPb0BJm1N1t7cCQ7ZM4,4892
 braintrust/test_otel.py,sha256=EyLOBMIZegvGyqqWl4vAlIazZfz1s4tWbgSAM68-2mM,31356
 braintrust/test_queue.py,sha256=872weDhTfxN1Vu4-cV8pZZ9XNE8UkX4WNcvc2mPtSDI,8431
 braintrust/test_score.py,sha256=PpfzNeaYhC4shMFqF85EczuMT0FCTr6l1rdEYWLe4g0,5857
@@ -50,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=EQlkRKEKEIegwawQp7zpbqvdiSpYZwNnHduewzJEXVk,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
@@ -108,11 +109,12 @@ 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
 braintrust/wrappers/test_utils.py,sha256=wpLNRSCAZ42eG4srJrGFMAlfBSkhaCOcn2QwJXkRQ7s,3005
+braintrust/wrappers/threads.py,sha256=rGQO6aFmFstrXA4xPEZCGTZJ1FC7CAS_4zF6ikszprc,3682
 braintrust/wrappers/agno/__init__.py,sha256=soVMMmm_nmMcX9rWxfmKw9ur03XNKRcDmlf3G4GbR4g,2306
 braintrust/wrappers/agno/agent.py,sha256=m3HCxQNotJviJswGYMxxsOnMilF-DNGqeFZFJa2Zhs4,6473
 braintrust/wrappers/agno/function_call.py,sha256=MSsHLGoBk98xGlm-CYnjE5pYBi-IJfokLhjtbYgJpX0,2206
@@ -123,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.2.dist-info/METADATA,sha256=lFY0wRN89qL-A6hoaipBAFXXs_m6dtspmP3ElXqvcAQ,3753
-braintrust-0.5.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-braintrust-0.5.2.dist-info/entry_points.txt,sha256=Zpc0_09g5xm8as5jHqqFq7fhwO0xHSNct_TrEMONS7Q,60
-braintrust-0.5.2.dist-info/top_level.txt,sha256=hw1-y-UFMf60RzAr8x_eM7SThbIuWfQsQIbVvqSF83A,11
-braintrust-0.5.2.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,,

braintrust/graph_util.py

@@ -1,147 +0,0 @@
-# Generic graph algorithms.
-
-import dataclasses
-from typing import Protocol
-
-
-# An UndirectedGraph consists of a set of vertex labels and a set of edges
-# between vertices.
-@dataclasses.dataclass
-class UndirectedGraph:
-    vertices: set[int]
-    edges: set[tuple[int, int]]
-
-
-# An AdjacencyListGraph is a mapping from vertex label to the list of vertices
-# where there is a directed edge from the key to the value.
-AdjacencyListGraph = dict[int, set[int]]
-
-
-class FirstVisitF(Protocol):
-    def __call__(self, vertex: int, *, parent_vertex: int | None, **kwargs) -> None:
-        """Extras:
-        - parent_vertex: the vertex which spawned the current vertex as its
-          child during the depth-first search. `parent_vertex` is guaranteed
-          to have been visited before the current one.
-        """
-        ...
-
-
-class LastVisitF(Protocol):
-    def __call__(self, vertex: int) -> None: ...
-
-
-def depth_first_search(
-    graph: AdjacencyListGraph,
-    first_visit_f: FirstVisitF | None = None,
-    last_visit_f: LastVisitF | None = None,
-    visitation_order: list[int] | None = None,
-) -> None:
-    """A general depth-first search algorithm over a directed graph. As it
-    traverses the graph, it invokes user-provided hooks when a vertex is *first*
-    visited (before visiting its children) and when it is *last* visited (after
-    visiting all its children).
-
-    The first_visit_f and last_visit_f functions may be passed additional
-    information beyond the vertex being visited as kwargs. See their type
-    signatures for more details. For future proofing, you will likely want to
-    capture **kwargs as a catchall in your functions.
-
-    An optional `visitation_order` can be specified, which controls the order in
-    which vertices will be first visited (outside of visiting them through a
-    different vertex). It can also be used to limit the set of starting vertices
-    considered. Otherwise, the DFS will visit all vertices in an unspecfied
-    order.
-    """
-
-    # Check the validity of the graph.
-    for vs in graph.values():
-        for v in vs:
-            assert v in graph
-
-    first_visited_vertices = set()
-    visitation_order = visitation_order if visitation_order is not None else graph.keys()
-    events = list(reversed([("first", x, dict(parent_vertex=None)) for x in visitation_order]))
-    while events:
-        event_type, vertex, extras = events.pop()
-
-        if event_type == "last":
-            if last_visit_f:
-                last_visit_f(vertex)
-            continue
-
-        # First visit of a node. If we've already visited it, skip.
-        if vertex in first_visited_vertices:
-            continue
-        first_visited_vertices.add(vertex)
-        if first_visit_f:
-            first_visit_f(vertex, parent_vertex=extras["parent_vertex"])
-
-        # Add 'first' visitation events for all the children of the vertex to
-        # the stack. But before this, add a 'last' visitation event for this
-        # vertex, so that once we've completed all the children, we get the last
-        # visitation event for this one.
-        events.append(("last", vertex, dict()))
-        for child in graph[vertex]:
-            events.append(("first", child, dict(parent_vertex=vertex)))
-
-
-def undirected_connected_components(graph: UndirectedGraph) -> list[list[int]]:
-    """Group together all the connected components of an undirected graph.
-    Return each group as a list of vertices.
-    """
-
-    # Perhaps the most performant way to implement this is via union find. But
-    # in lieu of that, we can use a depth-first search over a direct-ified
-    # version of the graph. Upon the first visit of each vertex, we assign it a
-    # label equal to the label of the parent vertex. If there is no parent
-    # vertex, we assign a new label. At the end, we can group together all the
-    # vertices with the same label.
-
-    directed_graph = {v: set() for v in graph.vertices}
-    for i, j in graph.edges:
-        directed_graph[i].add(j)
-        directed_graph[j].add(i)
-
-    label_counter = 0
-    vertex_labels = {}
-
-    def first_visit_f(vertex, parent_vertex, **kwargs):
-        if parent_vertex is not None:
-            label = vertex_labels[parent_vertex]
-        else:
-            nonlocal label_counter
-            # pylint: disable=used-before-assignment
-            label = label_counter
-            label_counter += 1
-        vertex_labels[vertex] = label
-
-    depth_first_search(directed_graph, first_visit_f=first_visit_f)
-    output = [[] for _ in range(label_counter)]
-    for vertex, label in vertex_labels.items():
-        output[label].append(vertex)
-
-    return output
-
-
-def topological_sort(graph: AdjacencyListGraph, visitation_order: list[int] | None = None) -> list[int]:
-    """The topological_sort function accepts a graph as input, with edges from
-    parents to children. It returns an ordering where parents are guaranteed to
-    come before their children.
-
-    The `visitation_order` is forwarded directly to `depth_first_search`.
-
-    Ordering with respect to cycles is unspecified. It is the caller's
-    responsibility to check for cycles if it matters.
-    """
-
-    # We use DFS, where upon the 'last' visitation of a node, we append it to
-    # the final ordering. Then we reverse the list at the end.
-    reverse_ordering = []
-
-    def last_visit_f(vertex, **kwargs):
-        reverse_ordering.append(vertex)
-
-    depth_first_search(graph, last_visit_f=last_visit_f, visitation_order=visitation_order)
-    reverse_ordering.reverse()
-    return reverse_ordering