feat(typesafe): Add integration with `system_one()` patches (#7716) · intertermux-code/sentry-python@cd3866d · GitHub
Skip to content

Commit cd3866d

Browse files
alexander-alderman-webbalexander-alderman-webb
authored
feat(typesafe): Add integration with system_one() patches (getsentry#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.
1 parent bd08efd commit cd3866d

13 files changed

Lines changed: 606 additions & 13 deletions

File tree

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

Lines changed: 4 additions & 0 deletions

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ typing = [
8181
"aiohttp>=3.14.3",
8282
"openai-agents>=0.22.0",
8383
"mistralai>=2.10.1",
84+
"typesafe-sdk>=0.7.2",
8485
]
8586
test = [
8687
"dataclasses ; python_full_version < '3.7'",

‎scripts/populate_tox/config.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,13 @@
561561
"package": "typer",
562562
"num_versions": 2,
563563
},
564+
"typesafe": {
565+
"package": "typesafe-sdk",
566+
"integration_name": "typesafe",
567+
"deps": {
568+
"*": ["pytest-asyncio"],
569+
},
570+
},
564571
"unleash": {
565572
"package": "UnleashClient",
566573
"num_versions": 2,

‎scripts/populate_tox/package_dependencies.jsonl‎

Lines changed: 5 additions & 8 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎scripts/populate_tox/releases.jsonl‎

Lines changed: 1 addition & 5 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎scripts/split_tox_gh_actions/split_tox_gh_actions.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
"mistral",
9696
"openai-base",
9797
"openai-notiktoken",
98+
"typesafe",
9899
],
99100
"Cloud": [
100101
"aws_lambda",

‎sentry_sdk/consts.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1238,6 +1238,7 @@ class OP:
12381238
GEN_AI_CHAT = "gen_ai.chat"
12391239
GEN_AI_CREATE_AGENT = "gen_ai.create_agent"
12401240
GEN_AI_EMBEDDINGS = "gen_ai.embeddings"
1241+
GEN_AI_EVALUATE = "gen_ai.evaluate"
12411242
GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool"
12421243
GEN_AI_TEXT_COMPLETION = "gen_ai.text_completion"
12431244
GEN_AI_HANDOFF = "gen_ai.handoff"

‎sentry_sdk/integrations/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ def iter_default_integrations(
168168
"strawberry": (0, 209, 5),
169169
"tornado": (6, 0),
170170
"typer": (0, 15),
171+
"typesafe": (0, 7),
171172
"unleash": (6, 0, 1),
172173
}
173174

‎sentry_sdk/integrations/typesafe.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
from functools import wraps
2+
from typing import TYPE_CHECKING
3+
4+
import sentry_sdk
5+
from sentry_sdk.ai.utils import get_start_span_function
6+
from sentry_sdk.consts import OP, SPANDATA
7+
from sentry_sdk.integrations import DidNotEnable, Integration
8+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
9+
10+
if TYPE_CHECKING:
11+
from typing import Any, Callable
12+
13+
try:
14+
from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient
15+
from typesafe_sdk._core.client.sync.client import TypeSafeClient
16+
except ImportError:
17+
raise DidNotEnable("typesafe-sdk not installed")
18+
19+
20+
class TypeSafeIntegration(Integration):
21+
identifier = "typesafe"
22+
origin = f"auto.ai.{identifier}"
23+
24+
@staticmethod
25+
def setup_once() -> None:
26+
TypeSafeClient.system_one = _wrap_system_one(TypeSafeClient.system_one) # type: ignore[method-assign]
27+
AsyncTypeSafeClient.system_one = _wrap_system_one_async( # type: ignore[method-assign]
28+
AsyncTypeSafeClient.system_one
29+
)
30+
31+
32+
def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]":
33+
@wraps(f)
34+
def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "Any":
35+
client = sentry_sdk.get_client()
36+
integration = client.get_integration(TypeSafeIntegration)
37+
if integration is None:
38+
return f(self, *args, **kwargs)
39+
40+
model = kwargs.get("model")
41+
if (
42+
model is None
43+
and hasattr(self, "_config")
44+
and hasattr(self._config, "default_model")
45+
):
46+
model = self._config.default_model
47+
48+
if has_span_streaming_enabled(client.options):
49+
span = sentry_sdk.traces.start_span(
50+
name=f"evaluate {model}".strip(),
51+
attributes={
52+
"sentry.op": OP.GEN_AI_EVALUATE,
53+
"sentry.origin": TypeSafeIntegration.origin,
54+
SPANDATA.GEN_AI_PROVIDER_NAME: "typesafe",
55+
SPANDATA.GEN_AI_OPERATION_NAME: "evaluate",
56+
},
57+
)
58+
set_on_span = span.set_attribute
59+
else:
60+
span = get_start_span_function()(
61+
op=OP.GEN_AI_EVALUATE,
62+
name=f"evaluate {model}".strip(),
63+
origin=TypeSafeIntegration.origin,
64+
)
65+
span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "typesafe")
66+
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "evaluate")
67+
set_on_span = span.set_data
68+
69+
with span:
70+
if model is not None:
71+
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
72+
73+
return f(self, *args, **kwargs)
74+
75+
return wrap_system_one
76+
77+
78+
def _wrap_system_one_async(f: "Callable[..., Any]") -> "Callable[..., Any]":
79+
@wraps(f)
80+
async def wrap_system_one_async(
81+
self: "AsyncTypeSafeClient", *args: "Any", **kwargs: "Any"
82+
) -> "Any":
83+
client = sentry_sdk.get_client()
84+
integration = client.get_integration(TypeSafeIntegration)
85+
if integration is None:
86+
return await f(self, *args, **kwargs)
87+
88+
model = kwargs.get("model")
89+
if (
90+
model is None
91+
and hasattr(self, "_config")
92+
and hasattr(self._config, "default_model")
93+
):
94+
model = self._config.default_model
95+
96+
if has_span_streaming_enabled(client.options):
97+
span = sentry_sdk.traces.start_span(
98+
name=f"evaluate {model}".strip(),
99+
attributes={
100+
"sentry.op": OP.GEN_AI_EVALUATE,
101+
"sentry.origin": TypeSafeIntegration.origin,
102+
SPANDATA.GEN_AI_PROVIDER_NAME: "typesafe",
103+
SPANDATA.GEN_AI_OPERATION_NAME: "evaluate",
104+
},
105+
)
106+
set_on_span = span.set_attribute
107+
else:
108+
span = get_start_span_function()(
109+
op=OP.GEN_AI_EVALUATE,
110+
name=f"evaluate {model}".strip(),
111+
origin=TypeSafeIntegration.origin,
112+
)
113+
span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "typesafe")
114+
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "evaluate")
115+
set_on_span = span.set_data
116+
117+
with span:
118+
if model is not None:
119+
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
120+
121+
return await f(self, *args, **kwargs)
122+
123+
return wrap_system_one_async

‎tests/integrations/typesafe/__init__.py‎

Lines changed: 3 additions & 0 deletions

0 commit comments

Comments
 (0)