Skip to content

Commit 7b028b6

Browse files
fix(integrations): allow explicit op parameter in ai_track (#4597)
fixes #4596 Co-authored-by: Daniel Szoke <[email protected]>
1 parent d32e2ee commit 7b028b6

File tree

2 files changed

+43
-2
lines changed

2 files changed

+43
-2
lines changed

sentry_sdk/ai/monitoring.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def decorator(f):
3232
def sync_wrapped(*args, **kwargs):
3333
# type: (Any, Any) -> Any
3434
curr_pipeline = _ai_pipeline_name.get()
35-
op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
35+
op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline")
3636

3737
with start_span(name=description, op=op, **span_kwargs) as span:
3838
for k, v in kwargs.pop("sentry_tags", {}).items():
@@ -61,7 +61,7 @@ def sync_wrapped(*args, **kwargs):
6161
async def async_wrapped(*args, **kwargs):
6262
# type: (Any, Any) -> Any
6363
curr_pipeline = _ai_pipeline_name.get()
64-
op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
64+
op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline")
6565

6666
with start_span(name=description, op=op, **span_kwargs) as span:
6767
for k, v in kwargs.pop("sentry_tags", {}).items():

tests/test_ai_monitoring.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,3 +119,44 @@ async def async_pipeline():
119119
assert ai_pipeline_span["tags"]["user"] == "czyber"
120120
assert ai_pipeline_span["data"]["some_data"] == "value"
121121
assert ai_run_span["description"] == "my async tool"
122+
123+
124+
def test_ai_track_with_explicit_op(sentry_init, capture_events):
125+
sentry_init(traces_sample_rate=1.0)
126+
events = capture_events()
127+
128+
@ai_track("my tool", op="custom.operation")
129+
def tool(**kwargs):
130+
pass
131+
132+
with sentry_sdk.start_transaction():
133+
tool()
134+
135+
transaction = events[0]
136+
assert transaction["type"] == "transaction"
137+
assert len(transaction["spans"]) == 1
138+
span = transaction["spans"][0]
139+
140+
assert span["description"] == "my tool"
141+
assert span["op"] == "custom.operation"
142+
143+
144+
@pytest.mark.asyncio
145+
async def test_ai_track_async_with_explicit_op(sentry_init, capture_events):
146+
sentry_init(traces_sample_rate=1.0)
147+
events = capture_events()
148+
149+
@ai_track("my async tool", op="custom.async.operation")
150+
async def async_tool(**kwargs):
151+
pass
152+
153+
with sentry_sdk.start_transaction():
154+
await async_tool()
155+
156+
transaction = events[0]
157+
assert transaction["type"] == "transaction"
158+
assert len(transaction["spans"]) == 1
159+
span = transaction["spans"][0]
160+
161+
assert span["description"] == "my async tool"
162+
assert span["op"] == "custom.async.operation"

0 commit comments

Comments
 (0)