1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
|
import pytest
import sentry_sdk
from sentry_sdk.ai.monitoring import ai_track
def test_ai_track(sentry_init, capture_events):
sentry_init(traces_sample_rate=1.0)
events = capture_events()
@ai_track("my tool")
def tool(**kwargs):
pass
@ai_track("some test pipeline")
def pipeline():
tool()
with sentry_sdk.start_transaction():
pipeline()
transaction = events[0]
assert transaction["type"] == "transaction"
assert len(transaction["spans"]) == 2
spans = transaction["spans"]
ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
assert ai_pipeline_span["description"] == "some test pipeline"
assert ai_run_span["description"] == "my tool"
def test_ai_track_with_tags(sentry_init, capture_events):
sentry_init(traces_sample_rate=1.0)
events = capture_events()
@ai_track("my tool")
def tool(**kwargs):
pass
@ai_track("some test pipeline")
def pipeline():
tool()
with sentry_sdk.start_transaction():
pipeline(sentry_tags={"user": "colin"}, sentry_data={"some_data": "value"})
transaction = events[0]
assert transaction["type"] == "transaction"
assert len(transaction["spans"]) == 2
spans = transaction["spans"]
ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
assert ai_pipeline_span["description"] == "some test pipeline"
print(ai_pipeline_span)
assert ai_pipeline_span["tags"]["user"] == "colin"
assert ai_pipeline_span["data"]["some_data"] == "value"
assert ai_run_span["description"] == "my tool"
@pytest.mark.asyncio
async def test_ai_track_async(sentry_init, capture_events):
sentry_init(traces_sample_rate=1.0)
events = capture_events()
@ai_track("my async tool")
async def async_tool(**kwargs):
pass
@ai_track("some async test pipeline")
async def async_pipeline():
await async_tool()
with sentry_sdk.start_transaction():
await async_pipeline()
transaction = events[0]
assert transaction["type"] == "transaction"
assert len(transaction["spans"]) == 2
spans = transaction["spans"]
ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
assert ai_pipeline_span["description"] == "some async test pipeline"
assert ai_run_span["description"] == "my async tool"
@pytest.mark.asyncio
async def test_ai_track_async_with_tags(sentry_init, capture_events):
sentry_init(traces_sample_rate=1.0)
events = capture_events()
@ai_track("my async tool")
async def async_tool(**kwargs):
pass
@ai_track("some async test pipeline")
async def async_pipeline():
await async_tool()
with sentry_sdk.start_transaction():
await async_pipeline(
sentry_tags={"user": "czyber"}, sentry_data={"some_data": "value"}
)
transaction = events[0]
assert transaction["type"] == "transaction"
assert len(transaction["spans"]) == 2
spans = transaction["spans"]
ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
assert ai_pipeline_span["description"] == "some async test pipeline"
assert ai_pipeline_span["tags"]["user"] == "czyber"
assert ai_pipeline_span["data"]["some_data"] == "value"
assert ai_run_span["description"] == "my async tool"
|
