feat(integrations): add TypeSafe instrumentation

py/noxfile.py

@@ -683,6 +683,18 @@ def test_mistral(session, version):
     _run_tests(session, f"{INTEGRATION_DIR}/mistral/test_mistral.py", version=version)
 
 
+TYPESAFE_VERSIONS = _get_matrix_versions("typesafe-sdk")
+
+
+@nox.session()
+@nox.parametrize("version", TYPESAFE_VERSIONS, ids=TYPESAFE_VERSIONS)
+def test_typesafe(session, version):
+    """Test the TypeSafe SDK integration."""
+    _install_test_deps(session)
+    _install_matrix_dep(session, "typesafe-sdk", version)
+    _run_tests(session, f"{INTEGRATION_DIR}/typesafe/test_typesafe.py", version=version)
+
+
 HUGGINGFACE_HUB_VERSIONS = _get_matrix_versions("huggingface-hub")
 
 

py/pyproject.toml

@@ -278,6 +278,7 @@ lint = [
     "huggingface-hub",
     "numpy",
     "strands-agents",
+    "typesafe-sdk",
     "temporalio",
     "pydantic-ai",
     "opentelemetry-instrumentation-openai==0.59.1",
@@ -498,6 +499,9 @@ latest = "openrouter==1.1.139"
 latest = "mistralai==2.10.0"
 "1.12.4" = "mistralai==1.12.4"
 
+[tool.braintrust.matrix.typesafe-sdk]
+latest = "typesafe-sdk==0.6.0"
+
 [tool.braintrust.matrix.huggingface-hub]
 # Floor pinned to 0.32.0: the earliest release that exposes the
 # ``provider="auto"`` routing mode the integration relies on for multi-
@@ -585,6 +589,7 @@ openrouter = ["openrouter"]
 pipecat = ["pipecat-ai"]
 pydantic_ai = ["pydantic-ai-integration", "pydantic-ai-wrap-openai"]
 strands = ["strands-agents"]
+typesafe = ["typesafe-sdk"]
 
 [tool.braintrust.vendor-packages]
 ai-sdk = "ai"
@@ -615,3 +620,4 @@ pipecat-ai = "pipecat"
 strands-agents = "strands"
 temporalio = "temporalio"
 transformers = "transformers"
+typesafe-sdk = "typesafe_sdk"

py/src/braintrust/auto.py

@@ -36,6 +36,7 @@
     StrandsIntegration,
     TemporalIntegration,
     TransformersIntegration,
+    TypeSafeIntegration,
 )
 from braintrust.integrations.base import BaseIntegration
 
@@ -86,6 +87,7 @@ def auto_instrument(
     temporal: bool = True,
     livekit_agents: bool = True,
     pipecat: bool = True,
+    typesafe: bool = True,
 ) -> dict[str, bool]:
     """
     Auto-instrument supported AI/ML libraries for Braintrust tracing.
@@ -125,6 +127,7 @@ def auto_instrument(
         temporal: Enable Temporal instrumentation (default: True)
         livekit_agents: Enable LiveKit Agents instrumentation (default: True)
         pipecat: Enable Pipecat AI instrumentation (default: True)
+        typesafe: Enable TypeSafe instrumentation (default: True)
 
     Returns:
         Dict mapping integration name to whether it was successfully instrumented.
@@ -228,6 +231,8 @@ def auto_instrument(
         results["livekit_agents"] = _instrument_integration(LiveKitAgentsIntegration)
     if pipecat:
         results["pipecat"] = _instrument_integration(PipecatIntegration)
+    if typesafe:
+        results["typesafe"] = _instrument_integration(TypeSafeIntegration)
 
     return results
 

py/src/braintrust/conftest.py

@@ -204,6 +204,7 @@ def setup_braintrust():
     os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy-api-key-for-vcr-tests")
     os.environ.setdefault("ANTHROPIC_API_KEY", "sk-ant-test-dummy-api-key-for-vcr-tests")
     os.environ.setdefault("MISTRAL_API_KEY", "mistral-test-dummy-api-key-for-vcr-tests")
+    os.environ.setdefault("TYPESAFE_API_KEY", "typesafe-test-dummy-api-key-for-vcr-tests")
     os.environ.setdefault("CO_API_KEY", os.getenv("COHERE_API_KEY", "co-test-dummy-api-key-for-vcr-tests"))
     os.environ.setdefault("COHERE_API_KEY", os.getenv("CO_API_KEY", "co-test-dummy-api-key-for-vcr-tests"))
 

py/src/braintrust/integrations/__init__.py

@@ -26,6 +26,7 @@
 from .strands import StrandsIntegration
 from .temporal import TemporalIntegration
 from .transformers import TransformersIntegration
+from .typesafe import TypeSafeIntegration
 
 
 __all__ = [
@@ -57,4 +58,5 @@
     "StrandsIntegration",
     "TemporalIntegration",
     "TransformersIntegration",
+    "TypeSafeIntegration",
 ]

py/src/braintrust/integrations/auto_test_scripts/test_auto_typesafe.py

@@ -0,0 +1,30 @@
+"""Test auto_instrument for TypeSafe."""
+
+import os
+
+from braintrust.auto import auto_instrument
+from braintrust.integrations.test_utils import autoinstrument_test_context
+from typesafe_sdk import Noul, TypeSafeClient
+
+
+results = auto_instrument()
+assert results.get("typesafe") is True
+assert auto_instrument().get("typesafe") is True
+
+with autoinstrument_test_context("test_auto_typesafe", integration="typesafe") as memory_logger:
+    with TypeSafeClient(api_key=os.environ["TYPESAFE_API_KEY"]) as client:
+        response = client.system_one(
+            state="The package arrived intact and on time.",
+            questions={"positive": Noul(instructions="Is this feedback positive?")},
+        )
+    assert 0 <= response.nouls["positive"].noul <= 1
+
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["metadata"]["provider"] == "typesafe"
+    assert span["metadata"]["model"].startswith("jev-")
+    assert span["input"]["questions"][0]["id"] == "positive"
+    assert span["output"]["answers"][0]["id"] == "positive"
+
+print("SUCCESS")

py/src/braintrust/integrations/typesafe/__init__.py

@@ -0,0 +1,26 @@
+"""Braintrust integration for the TypeSafe Python SDK."""
+
+from typing import Any
+
+from .integration import TypeSafeIntegration
+from .patchers import AsyncSystemOnePatcher, SystemOnePatcher
+
+
+def setup_typesafe() -> bool:
+    """Instrument installed TypeSafe sync and async clients."""
+    return TypeSafeIntegration.setup()
+
+
+def wrap_typesafe(client: Any) -> Any:
+    """Instrument a TypeSafe sync or async client in place."""
+    from typesafe_sdk import AsyncTypeSafeClient
+
+    patcher = AsyncSystemOnePatcher if isinstance(client, AsyncTypeSafeClient) else SystemOnePatcher
+    return patcher.wrap_target(client)
+
+
+__all__ = [
+    "TypeSafeIntegration",
+    "setup_typesafe",
+    "wrap_typesafe",
+]

py/src/braintrust/integrations/typesafe/integration.py

@@ -0,0 +1,15 @@
+"""TypeSafe integration orchestration."""
+
+from braintrust.integrations.base import BaseIntegration
+
+from .patchers import TypeSafePatcher
+
+
+class TypeSafeIntegration(BaseIntegration):
+    """Braintrust instrumentation for the TypeSafe Python SDK."""
+
+    name = "typesafe"
+    import_names = ("typesafe_sdk",)
+    distribution_names = ("typesafe-sdk",)
+    min_version = "0.6.0"
+    patchers = (TypeSafePatcher,)

py/src/braintrust/integrations/typesafe/patchers.py

@@ -0,0 +1,24 @@
+"""Patchers for TypeSafe sync and async System One calls."""
+
+from braintrust.integrations.base import CompositeFunctionWrapperPatcher, FunctionWrapperPatcher
+
+from .tracing import _async_system_one_wrapper, _system_one_wrapper
+
+
+class SystemOnePatcher(FunctionWrapperPatcher):
+    name = "typesafe.system_one"
+    target_module = "typesafe_sdk._core.client.sync.client"
+    target_path = "TypeSafeClient.system_one"
+    wrapper = _system_one_wrapper
+
+
+class AsyncSystemOnePatcher(FunctionWrapperPatcher):
+    name = "typesafe.async.system_one"
+    target_module = "typesafe_sdk._core.client.aio.client"
+    target_path = "AsyncTypeSafeClient.system_one"
+    wrapper = _async_system_one_wrapper
+
+
+class TypeSafePatcher(CompositeFunctionWrapperPatcher):
+    name = "typesafe.system_one.all"
+    sub_patchers = (SystemOnePatcher, AsyncSystemOnePatcher)