Skip to content

Commit e2eaaac

Browse files
shellmayrshellmayr
committed
simplify
1 parent 5c9cffb commit e2eaaac

File tree

5 files changed

+49
-704
lines changed

5 files changed

+49
-704
lines changed

‎sentry_sdk/ai/__init__.py‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,7 @@
1-
from .monitoring import record_token_usage # noqa: F401
21
from .utils import (
32
set_data_normalized,
43
GEN_AI_MESSAGE_ROLE_MAPPING,
54
GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING,
65
normalize_message_role,
76
normalize_message_roles,
87
) # noqa: F401
9-
10-
__all__ = [
11-
"record_token_usage",
12-
"set_data_normalized",
13-
"GEN_AI_MESSAGE_ROLE_MAPPING",
14-
"GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING",
15-
"normalize_message_role",
16-
"normalize_message_roles",
17-
]

‎sentry_sdk/integrations/openai.py‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -131,12 +131,7 @@ def _calculate_token_usage(
131131

132132
if hasattr(response, "usage"):
133133
input_tokens = _get_usage(response.usage, ["input_tokens", "prompt_tokens"])
134-
if hasattr(response.usage, "prompt_tokens_details"):
135-
input_tokens_cached = _get_usage(
136-
response.usage.prompt_tokens_details, ["cached_tokens"]
137-
)
138-
# OpenAI also supports input_tokens_details for compatibility
139-
elif hasattr(response.usage, "input_tokens_details"):
134+
if hasattr(response.usage, "input_tokens_details"):
140135
input_tokens_cached = _get_usage(
141136
response.usage.input_tokens_details, ["cached_tokens"]
142137
)
@@ -148,10 +143,6 @@ def _calculate_token_usage(
148143
output_tokens_reasoning = _get_usage(
149144
response.usage.output_tokens_details, ["reasoning_tokens"]
150145
)
151-
elif hasattr(response.usage, "completion_tokens_details"):
152-
output_tokens_reasoning = _get_usage(
153-
response.usage.completion_tokens_details, ["reasoning_tokens"]
154-
)
155146

156147
total_tokens = _get_usage(response.usage, ["total_tokens"])
157148

‎tests/integrations/anthropic/test_anthropic.py‎

Lines changed: 33 additions & 235 deletions
Original file line numberDiff line numberDiff line change
@@ -850,8 +850,10 @@ def test_collect_ai_data_with_input_json_delta():
850850
output_tokens = 20
851851
content_blocks = []
852852

853-
model, new_input_tokens, new_output_tokens, new_content_blocks = _collect_ai_data(
854-
event, model, input_tokens, output_tokens, content_blocks
853+
model, new_input_tokens, new_output_tokens, _, _, new_content_blocks = (
854+
_collect_ai_data(
855+
event, model, input_tokens, output_tokens, 0, 0, content_blocks
856+
)
855857
)
856858

857859
assert model is None
@@ -881,6 +883,8 @@ def test_set_output_data_with_input_json_delta(sentry_init):
881883
model="",
882884
input_tokens=10,
883885
output_tokens=20,
886+
cache_read_input_tokens=0,
887+
cache_write_input_tokens=0,
884888
content_blocks=[{"text": "".join(json_deltas), "type": "text"}],
885889
)
886890

@@ -1449,118 +1453,44 @@ def test_system_prompt_with_complex_structure(sentry_init, capture_events):
14491453

14501454

14511455
def test_cache_tokens_nonstreaming(sentry_init, capture_events):
1452-
"""Test that cache read and write tokens are properly tracked for non-streaming responses."""
1453-
sentry_init(
1454-
integrations=[AnthropicIntegration(include_prompts=True)],
1455-
traces_sample_rate=1.0,
1456-
send_default_pii=True,
1457-
)
1456+
"""Test cache read/write tokens are tracked for non-streaming responses."""
1457+
sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
14581458
events = capture_events()
14591459
client = Anthropic(api_key="z")
14601460

1461-
# Create a message with cache token usage
1462-
message_with_cache = Message(
1463-
id="id",
1464-
model="claude-3-5-sonnet-20241022",
1465-
role="assistant",
1466-
content=[TextBlock(type="text", text="Response using cache")],
1467-
type="message",
1468-
usage=Usage(
1469-
input_tokens=100,
1470-
output_tokens=50,
1471-
cache_read_input_tokens=80, # 80 tokens read from cache
1472-
cache_write_input_tokens=20, # 20 tokens written to cache
1473-
),
1474-
)
1475-
1476-
client.messages._post = mock.Mock(return_value=message_with_cache)
1477-
1478-
messages = [{"role": "user", "content": "Hello"}]
1479-
1480-
with start_transaction(name="anthropic"):
1481-
response = client.messages.create(
1482-
max_tokens=1024, messages=messages, model="claude-3-5-sonnet-20241022"
1461+
client.messages._post = mock.Mock(
1462+
return_value=Message(
1463+
id="id",
1464+
model="claude-3-5-sonnet-20241022",
1465+
role="assistant",
1466+
content=[TextBlock(type="text", text="Response")],
1467+
type="message",
1468+
usage=Usage(
1469+
input_tokens=100,
1470+
output_tokens=50,
1471+
cache_read_input_tokens=80,
1472+
cache_creation_input_tokens=20,
1473+
),
14831474
)
1484-
1485-
assert response == message_with_cache
1486-
usage = response.usage
1487-
1488-
assert usage.input_tokens == 100
1489-
assert usage.output_tokens == 50
1490-
assert usage.cache_read_input_tokens == 80
1491-
assert usage.cache_write_input_tokens == 20
1492-
1493-
assert len(events) == 1
1494-
(event,) = events
1495-
1496-
assert event["type"] == "transaction"
1497-
assert len(event["spans"]) == 1
1498-
(span,) = event["spans"]
1499-
1500-
assert span["op"] == OP.GEN_AI_CHAT
1501-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1502-
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 50
1503-
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 150
1504-
# Check cache-related tokens
1505-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
1506-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
1507-
1508-
1509-
def test_cache_tokens_only_reads(sentry_init, capture_events):
1510-
"""Test tracking when only cache reads are present (no writes)."""
1511-
sentry_init(
1512-
integrations=[AnthropicIntegration(include_prompts=True)],
1513-
traces_sample_rate=1.0,
1514-
send_default_pii=True,
1515-
)
1516-
events = capture_events()
1517-
client = Anthropic(api_key="z")
1518-
1519-
# Message with only cache reads, no writes
1520-
message_cache_read_only = Message(
1521-
id="id",
1522-
model="claude-3-5-sonnet-20241022",
1523-
role="assistant",
1524-
content=[TextBlock(type="text", text="Response")],
1525-
type="message",
1526-
usage=Usage(
1527-
input_tokens=100,
1528-
output_tokens=50,
1529-
cache_read_input_tokens=100, # All tokens read from cache
1530-
cache_write_input_tokens=0, # No new cache writes
1531-
),
15321475
)
15331476

1534-
client.messages._post = mock.Mock(return_value=message_cache_read_only)
1535-
15361477
with start_transaction(name="anthropic"):
15371478
client.messages.create(
15381479
max_tokens=1024,
15391480
messages=[{"role": "user", "content": "Hello"}],
15401481
model="claude-3-5-sonnet-20241022",
15411482
)
15421483

1543-
assert len(events) == 1
1544-
(event,) = events
1545-
(span,) = event["spans"]
1546-
1547-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 100
1548-
# Cache write should not be present when it's 0
1549-
assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE not in span["data"]
1484+
(span,) = events[0]["spans"]
1485+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
1486+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
15501487

15511488

15521489
def test_cache_tokens_streaming(sentry_init, capture_events):
1553-
"""Test that cache tokens are tracked correctly for streaming responses."""
1554-
sentry_init(
1555-
integrations=[AnthropicIntegration(include_prompts=True)],
1556-
traces_sample_rate=1.0,
1557-
send_default_pii=True,
1558-
)
1559-
events = capture_events()
1490+
"""Test cache tokens are tracked for streaming responses."""
15601491
client = Anthropic(api_key="z")
1561-
1562-
# Create streaming events with cache usage
1563-
stream_events = [
1492+
returned_stream = Stream(cast_to=None, response=None, client=client)
1493+
returned_stream._iterator = [
15641494
MessageStartEvent(
15651495
type="message_start",
15661496
message=Message(
@@ -1573,162 +1503,30 @@ def test_cache_tokens_streaming(sentry_init, capture_events):
15731503
input_tokens=100,
15741504
output_tokens=0,
15751505
cache_read_input_tokens=80,
1576-
cache_write_input_tokens=20,
1506+
cache_creation_input_tokens=20,
15771507
),
15781508
),
15791509
),
1580-
ContentBlockDeltaEvent(
1581-
type="content_block_delta",
1582-
index=0,
1583-
delta=TextDelta(type="text_delta", text="Hello"),
1584-
),
15851510
MessageDeltaEvent(
15861511
type="message_delta",
15871512
delta=Delta(stop_reason="end_turn"),
15881513
usage=MessageDeltaUsage(output_tokens=10),
15891514
),
15901515
]
15911516

1592-
mock_stream = mock.MagicMock(spec=Stream)
1593-
mock_stream.__iter__ = mock.Mock(return_value=iter(stream_events))
1594-
mock_stream._iterator = iter(stream_events)
1595-
1596-
client.messages._post = mock.Mock(return_value=mock_stream)
1597-
1598-
with start_transaction(name="anthropic"):
1599-
stream = client.messages.create(
1600-
max_tokens=1024,
1601-
messages=[{"role": "user", "content": "Hello"}],
1602-
model="claude-3-5-sonnet-20241022",
1603-
stream=True,
1604-
)
1605-
# Consume the stream
1606-
for _ in stream:
1607-
pass
1608-
1609-
assert len(events) == 1
1610-
(event,) = events
1611-
(span,) = event["spans"]
1612-
1613-
assert span["op"] == OP.GEN_AI_CHAT
1614-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1615-
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10
1616-
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 110
1617-
# Check streaming cache tokens
1618-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
1619-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
1620-
assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True
1621-
1622-
1623-
@pytest.mark.asyncio
1624-
async def test_cache_tokens_streaming_async(sentry_init, capture_events):
1625-
"""Test that cache tokens are tracked correctly for async streaming responses."""
1626-
sentry_init(
1627-
integrations=[AnthropicIntegration(include_prompts=True)],
1628-
traces_sample_rate=1.0,
1629-
send_default_pii=True,
1630-
)
1517+
sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
16311518
events = capture_events()
1632-
client = AsyncAnthropic(api_key="z")
1633-
1634-
async def async_iterator(values):
1635-
for value in values:
1636-
yield value
1637-
1638-
# Create streaming events with cache usage
1639-
stream_events = [
1640-
MessageStartEvent(
1641-
type="message_start",
1642-
message=Message(
1643-
id="id",
1644-
model="claude-3-5-sonnet-20241022",
1645-
role="assistant",
1646-
content=[],
1647-
type="message",
1648-
usage=Usage(
1649-
input_tokens=100,
1650-
output_tokens=0,
1651-
cache_read_input_tokens=80,
1652-
cache_write_input_tokens=20,
1653-
),
1654-
),
1655-
),
1656-
ContentBlockDeltaEvent(
1657-
type="content_block_delta",
1658-
index=0,
1659-
delta=TextDelta(type="text_delta", text="Hello"),
1660-
),
1661-
MessageDeltaEvent(
1662-
type="message_delta",
1663-
delta=Delta(stop_reason="end_turn"),
1664-
usage=MessageDeltaUsage(output_tokens=10),
1665-
),
1666-
]
1667-
1668-
mock_stream = mock.MagicMock(spec=AsyncStream)
1669-
mock_stream.__aiter__ = mock.Mock(return_value=async_iterator(stream_events))
1670-
mock_stream._iterator = async_iterator(stream_events)
1671-
1672-
client.messages._post = mock.Mock(return_value=mock_stream)
1519+
client.messages._post = mock.Mock(return_value=returned_stream)
16731520

16741521
with start_transaction(name="anthropic"):
1675-
stream = await client.messages.create(
1522+
for _ in client.messages.create(
16761523
max_tokens=1024,
16771524
messages=[{"role": "user", "content": "Hello"}],
16781525
model="claude-3-5-sonnet-20241022",
16791526
stream=True,
1680-
)
1681-
# Consume the stream
1682-
async for _ in stream:
1527+
):
16831528
pass
16841529

1685-
assert len(events) == 1
1686-
(event,) = events
1687-
(span,) = event["spans"]
1688-
1689-
assert span["op"] == OP.GEN_AI_CHAT
1690-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1691-
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10
1692-
# Check async streaming cache tokens
1530+
(span,) = events[0]["spans"]
16931531
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
16941532
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
1695-
1696-
1697-
def test_no_cache_tokens(sentry_init, capture_events):
1698-
"""Test that requests without cache usage don't have cache fields."""
1699-
sentry_init(
1700-
integrations=[AnthropicIntegration(include_prompts=True)],
1701-
traces_sample_rate=1.0,
1702-
send_default_pii=True,
1703-
)
1704-
events = capture_events()
1705-
client = Anthropic(api_key="z")
1706-
1707-
# Message without any cache usage
1708-
message_no_cache = Message(
1709-
id="id",
1710-
model="claude-3-5-sonnet-20241022",
1711-
role="assistant",
1712-
content=[TextBlock(type="text", text="Response")],
1713-
type="message",
1714-
usage=Usage(input_tokens=100, output_tokens=50),
1715-
)
1716-
1717-
client.messages._post = mock.Mock(return_value=message_no_cache)
1718-
1719-
with start_transaction(name="anthropic"):
1720-
client.messages.create(
1721-
max_tokens=1024,
1722-
messages=[{"role": "user", "content": "Hello"}],
1723-
model="claude-3-5-sonnet-20241022",
1724-
)
1725-
1726-
assert len(events) == 1
1727-
(event,) = events
1728-
(span,) = event["spans"]
1729-
1730-
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 100
1731-
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 50
1732-
# Cache fields should not be present
1733-
assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED not in span["data"]
1734-
assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE not in span["data"]

0 commit comments

Comments
 (0)