feat(typesafe): Add integration with system_one() patches (#7716)

Add a `typesafe_sdk` integration, the `typesafe` tox environment, and `typesafe_sdk` to the uv typing group.

Set the `gen_ai.provider.name`, `gen_ai.operation.name` and `gen_ai.request.model` attributes.

Author: alexander-alderman-webb

.github/workflows/test-integrations-ai.yml

@@ -69,6 +69,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh "py${{ matrix.python-version }}-openai-notiktoken"
+      - name: Test typesafe
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-typesafe"
       - name: Upload coverage data
         if: ${{ !cancelled() }}
         continue-on-error: true

pyproject.toml

@@ -81,6 +81,7 @@ typing = [
     "aiohttp>=3.14.3",
     "openai-agents>=0.22.0",
     "mistralai>=2.10.1",
+    "typesafe-sdk>=0.7.2",
 ]
 test = [
     "dataclasses ; python_full_version < '3.7'",

scripts/populate_tox/config.py

@@ -561,6 +561,13 @@
         "package": "typer",
         "num_versions": 2,
     },
+    "typesafe": {
+        "package": "typesafe-sdk",
+        "integration_name": "typesafe",
+        "deps": {
+            "*": ["pytest-asyncio"],
+        },
+    },
     "unleash": {
         "package": "UnleashClient",
         "num_versions": 2,

scripts/populate_tox/package_dependencies.jsonl

@@ -95,6 +95,7 @@
         "mistral",
         "openai-base",
         "openai-notiktoken",
+        "typesafe",
     ],
     "Cloud": [
         "aws_lambda",

scripts/populate_tox/releases.jsonl

@@ -1238,6 +1238,7 @@ class OP:
     GEN_AI_CHAT = "gen_ai.chat"
     GEN_AI_CREATE_AGENT = "gen_ai.create_agent"
     GEN_AI_EMBEDDINGS = "gen_ai.embeddings"
+    GEN_AI_EVALUATE = "gen_ai.evaluate"
     GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool"
     GEN_AI_TEXT_COMPLETION = "gen_ai.text_completion"
     GEN_AI_HANDOFF = "gen_ai.handoff"

scripts/split_tox_gh_actions/split_tox_gh_actions.py

@@ -168,6 +168,7 @@ def iter_default_integrations(
     "strawberry": (0, 209, 5),
     "tornado": (6, 0),
     "typer": (0, 15),
+    "typesafe": (0, 7),
     "unleash": (6, 0, 1),
 }
 

sentry_sdk/consts.py

@@ -0,0 +1,123 @@
+from functools import wraps
+from typing import TYPE_CHECKING
+
+import sentry_sdk
+from sentry_sdk.ai.utils import get_start_span_function
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+try:
+    from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient
+    from typesafe_sdk._core.client.sync.client import TypeSafeClient
+except ImportError:
+    raise DidNotEnable("typesafe-sdk not installed")
+
+
+class TypeSafeIntegration(Integration):
+    identifier = "typesafe"
+    origin = f"auto.ai.{identifier}"
+
+    @staticmethod
+    def setup_once() -> None:
+        TypeSafeClient.system_one = _wrap_system_one(TypeSafeClient.system_one)  # type: ignore[method-assign]
+        AsyncTypeSafeClient.system_one = _wrap_system_one_async(  # type: ignore[method-assign]
+            AsyncTypeSafeClient.system_one
+        )
+
+
+def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]":
+    @wraps(f)
+    def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(TypeSafeIntegration)
+        if integration is None:
+            return f(self, *args, **kwargs)
+
+        model = kwargs.get("model")
+        if (
+            model is None
+            and hasattr(self, "_config")
+            and hasattr(self._config, "default_model")
+        ):
+            model = self._config.default_model
+
+        if has_span_streaming_enabled(client.options):
+            span = sentry_sdk.traces.start_span(
+                name=f"evaluate {model}".strip(),
+                attributes={
+                    "sentry.op": OP.GEN_AI_EVALUATE,
+                    "sentry.origin": TypeSafeIntegration.origin,
+                    SPANDATA.GEN_AI_PROVIDER_NAME: "typesafe",
+                    SPANDATA.GEN_AI_OPERATION_NAME: "evaluate",
+                },
+            )
+            set_on_span = span.set_attribute
+        else:
+            span = get_start_span_function()(
+                op=OP.GEN_AI_EVALUATE,
+                name=f"evaluate {model}".strip(),
+                origin=TypeSafeIntegration.origin,
+            )
+            span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "typesafe")
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "evaluate")
+            set_on_span = span.set_data
+
+        with span:
+            if model is not None:
+                set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+
+            return f(self, *args, **kwargs)
+
+    return wrap_system_one
+
+
+def _wrap_system_one_async(f: "Callable[..., Any]") -> "Callable[..., Any]":
+    @wraps(f)
+    async def wrap_system_one_async(
+        self: "AsyncTypeSafeClient", *args: "Any", **kwargs: "Any"
+    ) -> "Any":
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(TypeSafeIntegration)
+        if integration is None:
+            return await f(self, *args, **kwargs)
+
+        model = kwargs.get("model")
+        if (
+            model is None
+            and hasattr(self, "_config")
+            and hasattr(self._config, "default_model")
+        ):
+            model = self._config.default_model
+
+        if has_span_streaming_enabled(client.options):
+            span = sentry_sdk.traces.start_span(
+                name=f"evaluate {model}".strip(),
+                attributes={
+                    "sentry.op": OP.GEN_AI_EVALUATE,
+                    "sentry.origin": TypeSafeIntegration.origin,
+                    SPANDATA.GEN_AI_PROVIDER_NAME: "typesafe",
+                    SPANDATA.GEN_AI_OPERATION_NAME: "evaluate",
+                },
+            )
+            set_on_span = span.set_attribute
+        else:
+            span = get_start_span_function()(
+                op=OP.GEN_AI_EVALUATE,
+                name=f"evaluate {model}".strip(),
+                origin=TypeSafeIntegration.origin,
+            )
+            span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "typesafe")
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "evaluate")
+            set_on_span = span.set_data
+
+        with span:
+            if model is not None:
+                set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+
+            return await f(self, *args, **kwargs)
+
+    return wrap_system_one_async

sentry_sdk/integrations/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("typesafe_sdk")