Skip to content

Commit dadbc79

Browse files
shellmayrshellmayr
committed
clarify
1 parent 33c63f6 commit dadbc79

6 files changed

Lines changed: 198 additions & 186 deletions

File tree

‎sentry_sdk/ai/span_config.py‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import sentry_sdk
22
from sentry_sdk.consts import SPANDATA
3+
from sentry_sdk.ai.monitoring import record_token_usage
34
from sentry_sdk.ai.utils import (
5+
get_first_from_sources,
46
set_data_normalized,
7+
set_span_data_from_sources,
58
normalize_message_roles,
69
truncate_and_annotate_messages,
710
)
@@ -10,11 +13,11 @@
1013
from typing import TYPE_CHECKING
1114

1215
if TYPE_CHECKING:
13-
from typing import Any, Dict
16+
from typing import Any, Dict, List, Optional
1417
from sentry_sdk.tracing import Span
1518

1619

17-
def set_input_span_data(span, kwargs, integration, config, span_data=None):
20+
def set_request_span_data(span, kwargs, integration, config, span_data=None):
1821
# type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None
1922
"""
2023
Set input span data from a declarative config.
@@ -57,3 +60,49 @@ def set_input_span_data(span, kwargs, integration, config, span_data=None):
5760
if kwarg_key in kwargs:
5861
value = kwargs[kwarg_key]
5962
set_data_normalized(span, span_attr, value)
63+
64+
65+
def set_response_span_data(span, response, include_pii, response_config, collected_text=None):
66+
# type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None
67+
"""
68+
Set response span data from a declarative config.
69+
70+
response_config keys:
71+
sources: dict - always set from response object
72+
pii_sources: dict - only when PII allowed
73+
extract_text: (response) -> list[str] | None (PII only)
74+
usage: dict with input_tokens/output_tokens source paths
75+
collected_text: pre-collected streaming text (overrides extract_text)
76+
"""
77+
set_span_data_from_sources(
78+
span, response, response_config.get("sources", {}), require_truthy=False
79+
)
80+
81+
if include_pii:
82+
pii_sources = response_config.get("pii_sources")
83+
if pii_sources:
84+
set_span_data_from_sources(
85+
span, response, pii_sources, require_truthy=True
86+
)
87+
if collected_text:
88+
set_data_normalized(
89+
span, SPANDATA.GEN_AI_RESPONSE_TEXT, ["".join(collected_text)]
90+
)
91+
else:
92+
extract_text = response_config.get("extract_text")
93+
if extract_text:
94+
texts = extract_text(response)
95+
if texts:
96+
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts)
97+
98+
usage_config = response_config.get("usage")
99+
if usage_config:
100+
record_token_usage(
101+
span,
102+
input_tokens=get_first_from_sources(
103+
response, usage_config.get("input_tokens", [])
104+
),
105+
output_tokens=get_first_from_sources(
106+
response, usage_config.get("output_tokens", [])
107+
),
108+
)

‎sentry_sdk/ai/utils.py‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
99

1010
if TYPE_CHECKING:
11-
from typing import Any, Callable, Dict, List, Optional, Tuple
11+
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
1212

1313
from sentry_sdk.tracing import Span
1414

@@ -725,3 +725,32 @@ def set_conversation_id(conversation_id: str) -> None:
725725
"""
726726
scope = sentry_sdk.get_current_scope()
727727
scope.set_conversation_id(conversation_id)
728+
729+
730+
def transitive_getattr(obj, *attrs):
731+
# type: (Any, str) -> Any
732+
current = obj
733+
for attr in attrs:
734+
current = getattr(current, attr, None)
735+
if current is None:
736+
return None
737+
return current
738+
739+
740+
def get_first_from_sources(obj, source_paths, require_truthy=False):
741+
# type: (Any, Sequence[tuple[str, ...]], bool) -> Any
742+
for source_path in source_paths:
743+
value = transitive_getattr(obj, *source_path)
744+
if not value:
745+
continue
746+
if not require_truthy or value:
747+
return value
748+
return None
749+
750+
751+
def set_span_data_from_sources(span, obj, target_sources, require_truthy):
752+
# type: (Any, Any, Mapping[str, Sequence[tuple[str, ...]]], bool) -> None
753+
for spandata_key, source_paths in target_sources.items():
754+
value = get_first_from_sources(obj, source_paths, require_truthy=require_truthy)
755+
if value is not None:
756+
set_data_normalized(span, spandata_key, value)

‎sentry_sdk/integrations/cohere/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from sentry_sdk.ai.monitoring import record_token_usage
55
from sentry_sdk.consts import OP, SPANDATA
6-
from sentry_sdk.ai.span_config import set_input_span_data
6+
from sentry_sdk.ai.span_config import set_request_span_data
77

88
from typing import TYPE_CHECKING
99

@@ -92,7 +92,7 @@ def new_embed(*args, **kwargs):
9292
name=f"embeddings {model}".strip(),
9393
origin=CohereIntegration.origin,
9494
) as span:
95-
set_input_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG)
95+
set_request_span_data(span, kwargs, integration, COHERE_EMBED_CONFIG)
9696

9797
try:
9898
res = f(*args, **kwargs)

‎sentry_sdk/integrations/cohere/utils.py‎

Lines changed: 0 additions & 35 deletions
This file was deleted.

‎sentry_sdk/integrations/cohere/v1.py‎

Lines changed: 47 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
import sys
22
from functools import wraps
33

4-
from sentry_sdk.ai.monitoring import record_token_usage
5-
from sentry_sdk.ai.span_config import set_input_span_data
6-
from sentry_sdk.ai.utils import set_data_normalized, transform_message_content
4+
from sentry_sdk.ai.span_config import set_request_span_data, set_response_span_data
5+
from sentry_sdk.ai.utils import (
6+
get_first_from_sources,
7+
transform_message_content,
8+
)
79
from sentry_sdk.consts import OP, SPANDATA
810

911
from typing import TYPE_CHECKING
@@ -17,10 +19,6 @@
1719
CohereIntegration,
1820
_capture_exception,
1921
)
20-
from sentry_sdk.integrations.cohere.utils import (
21-
get_first_from_sources,
22-
set_span_data_from_sources,
23-
)
2422
from sentry_sdk.scope import should_send_default_pii
2523
from sentry_sdk.utils import capture_internal_exceptions, reraise
2624

@@ -38,34 +36,41 @@
3836
except ImportError:
3937
_has_chat_types = False
4038

39+
def _extract_response_text(response):
40+
# type: (Any) -> list[str] | None
41+
text = getattr(response, "text", None)
42+
return [text] if text is not None else None
43+
44+
4145
COHERE_V1_CHAT_CONFIG = {
4246
"static": {
4347
SPANDATA.GEN_AI_SYSTEM: "cohere",
4448
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
4549
},
4650
"extract_messages": lambda kw: _extract_messages(kw),
51+
"response": {
52+
"sources": {
53+
SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)],
54+
SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)],
55+
},
56+
"pii_sources": {
57+
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)],
58+
},
59+
"extract_text": _extract_response_text,
60+
"usage": {
61+
"input_tokens": [
62+
("meta", "billed_units", "input_tokens"),
63+
("meta", "tokens", "input_tokens"),
64+
],
65+
"output_tokens": [
66+
("meta", "billed_units", "output_tokens"),
67+
("meta", "tokens", "output_tokens"),
68+
],
69+
},
70+
},
71+
"stream_response_object": [("response",)],
4772
}
4873

49-
CHAT_RESPONSE_SOURCES = {
50-
SPANDATA.GEN_AI_RESPONSE_ID: [("generation_id",)],
51-
SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: [("finish_reason",)],
52-
}
53-
PII_CHAT_RESPONSE_SOURCES = {
54-
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("tool_calls",)],
55-
}
56-
CHAT_RESPONSE_TEXT_SOURCES = [("text",)]
57-
CHAT_USAGE_TOKEN_SOURCES = {
58-
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: [
59-
("meta", "billed_units", "input_tokens"),
60-
("meta", "tokens", "input_tokens"),
61-
],
62-
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: [
63-
("meta", "billed_units", "output_tokens"),
64-
("meta", "tokens", "output_tokens"),
65-
],
66-
}
67-
STREAM_RESPONSE_SOURCES = [("response",)]
68-
6974

7075
def setup_v1(wrap_embed_fn):
7176
# type: (Callable[..., Any]) -> None
@@ -113,19 +118,20 @@ def new_chat(*args, **kwargs):
113118
reraise(*exc_info)
114119

115120
with capture_internal_exceptions():
116-
span_data = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}
117-
if model:
118-
span_data[SPANDATA.GEN_AI_REQUEST_MODEL] = model
119-
set_input_span_data(
121+
span_data = {
122+
SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming,
123+
SPANDATA.GEN_AI_REQUEST_MODEL: model if model else None,
124+
}
125+
set_request_span_data(
120126
span, kwargs, integration, COHERE_V1_CHAT_CONFIG, span_data
121127
)
122128

123129
if streaming:
124130
return _iter_stream_events(res, span, include_pii)
125-
if isinstance(res, NonStreamedChatResponse):
126-
_collect_response_fields(span, res, include_pii=include_pii)
127131
else:
128-
set_data_normalized(span, "unknown_response", True)
132+
set_response_span_data(
133+
span, res, include_pii, COHERE_V1_CHAT_CONFIG["response"]
134+
)
129135
return res
130136

131137
return new_chat
@@ -154,36 +160,11 @@ def _iter_stream_events(old_iterator, span, include_pii):
154160
if isinstance(x, ChatStreamEndEvent) or isinstance(
155161
x, StreamEndStreamedChatResponse
156162
):
157-
_collect_v1_stream_end_fields(span, x, include_pii)
163+
response = get_first_from_sources(
164+
x, COHERE_V1_CHAT_CONFIG["stream_response_object"]
165+
)
166+
if response is not None:
167+
set_response_span_data(
168+
span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"]
169+
)
158170
yield x
159-
160-
161-
def _collect_v1_stream_end_fields(span, event, include_pii):
162-
# type: (Any, Any, bool) -> None
163-
response = get_first_from_sources(event, STREAM_RESPONSE_SOURCES)
164-
if response is not None:
165-
_collect_response_fields(span, response, include_pii)
166-
167-
168-
def _collect_response_fields(span, response, include_pii):
169-
# type: (Any, Any, bool) -> None
170-
if include_pii:
171-
text = get_first_from_sources(response, CHAT_RESPONSE_TEXT_SOURCES)
172-
if text is not None:
173-
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, [text])
174-
set_span_data_from_sources(
175-
span, response, PII_CHAT_RESPONSE_SOURCES, require_truthy=False
176-
)
177-
178-
set_span_data_from_sources(
179-
span, response, CHAT_RESPONSE_SOURCES, require_truthy=False
180-
)
181-
record_token_usage(
182-
span,
183-
input_tokens=get_first_from_sources(
184-
response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS]
185-
),
186-
output_tokens=get_first_from_sources(
187-
response, CHAT_USAGE_TOKEN_SOURCES[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS]
188-
),
189-
)

0 commit comments

Comments
 (0)