Skip to content

Commit 3e077fe

Browse files
shellmayrshellmayr
committed
typing
1 parent 43bad0c commit 3e077fe

4 files changed

Lines changed: 103 additions & 70 deletions

File tree

‎sentry_sdk/ai/span_config.py‎

Lines changed: 54 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -13,24 +13,64 @@
1313
from typing import TYPE_CHECKING
1414

1515
if TYPE_CHECKING:
16-
from typing import Any, Dict, List, Optional
16+
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
17+
from typing_extensions import TypedDict
18+
1719
from sentry_sdk.tracing import Span
1820

21+
# Source paths: list of attribute chains to try in order.
22+
# e.g. [("meta", "billed_units", "input_tokens"), ("meta", "tokens", "input_tokens")]
23+
SourcePaths = Sequence[Tuple[str, ...]]
24+
25+
# Maps a SPANDATA key to source paths on the response object.
26+
# e.g. {SPANDATA.GEN_AI_RESPONSE_ID: [("id",)]}
27+
SourceMapping = Dict[str, SourcePaths]
28+
29+
class UsageConfig(TypedDict, total=False):
30+
"""Declarative token usage extraction paths (from response object)."""
31+
32+
input_tokens: SourcePaths
33+
output_tokens: SourcePaths
34+
total_tokens: SourcePaths
35+
36+
class ResponseConfig(TypedDict, total=False):
37+
"""Declarative response span data config."""
38+
39+
# Attributes always extracted from the response object.
40+
sources: SourceMapping
41+
# Attributes extracted only when PII sending is enabled.
42+
pii_sources: SourceMapping
43+
# Custom extractor for response text (PII only).
44+
# Returns list of text strings, or None.
45+
extract_text: Callable[[Any], Optional[List[str]]]
46+
# Declarative token usage paths.
47+
usage: UsageConfig
48+
49+
class OperationConfig(TypedDict, total=False):
50+
"""Full declarative config for an AI operation (chat, embeddings, etc.)."""
51+
52+
# Key/value pairs set on every span unconditionally.
53+
static: Dict[str, Any]
54+
# Maps kwarg names to SPANDATA keys (always set if present in kwargs).
55+
params: Dict[str, str]
56+
# Maps kwarg names to SPANDATA keys (only set when PII is enabled).
57+
pii_params: Dict[str, str]
58+
# Extracts messages from kwargs for the span.
59+
extract_messages: Callable[[Dict[str, Any]], Optional[List[Dict[str, Any]]]]
60+
# SPANDATA key for messages (default: GEN_AI_REQUEST_MESSAGES).
61+
message_target: str
62+
# Non-streaming response config.
63+
response: ResponseConfig
64+
# Streaming response config (different attribute paths).
65+
stream_response: ResponseConfig
66+
# Source paths to extract a full response object from a stream-end event
67+
# (V1 pattern: reuse "response" config after extracting).
68+
stream_response_object: SourcePaths
69+
1970

2071
def set_request_span_data(span, kwargs, integration, config, span_data=None):
2172
# type: (Span, Dict[str, Any], Any, Dict[str, Any], Dict[str, Any] | None) -> None
22-
"""
23-
Set input span data from a declarative config.
24-
25-
Config keys:
26-
static: dict - key/value pairs to set unconditionally
27-
params: dict - kwargs key -> span attr (always set if present)
28-
pii_params: dict - kwargs key -> span attr (only when PII allowed)
29-
extract_messages: callable(kwargs) -> list or None
30-
message_target: str - span attr for messages (default: GEN_AI_REQUEST_MESSAGES)
31-
32-
span_data: additional key/value pairs for dynamic per-call values
33-
"""
73+
"""Set request/static span data from a declarative config."""
3474
for key, value in config.get("static", {}).items():
3575
set_data_normalized(span, key, value)
3676
if span_data:
@@ -66,16 +106,7 @@ def set_response_span_data(
66106
span, response, include_pii, response_config, collected_text=None
67107
):
68108
# type: (Span, Any, bool, Dict[str, Any], Optional[List[str]]) -> None
69-
"""
70-
Set response span data from a declarative config.
71-
72-
response_config keys:
73-
sources: dict - always set from response object
74-
pii_sources: dict - only when PII allowed
75-
extract_text: (response) -> list[str] | None (PII only)
76-
usage: dict with input_tokens/output_tokens source paths
77-
collected_text: pre-collected streaming text (overrides extract_text)
78-
"""
109+
"""Set response span data from a declarative config."""
79110
set_span_data_from_sources(
80111
span, response, response_config.get("sources", {}), require_truthy=False
81112
)

‎sentry_sdk/integrations/cohere/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
if TYPE_CHECKING:
1212
from typing import Any, Callable
13+
from sentry_sdk.ai.span_config import OperationConfig
1314

1415
import sentry_sdk
1516
from sentry_sdk.integrations import DidNotEnable, Integration
@@ -30,7 +31,7 @@ def _normalize_embedding_input(texts):
3031
return [texts]
3132

3233

33-
COHERE_EMBED_CONFIG = {
34+
COHERE_EMBED_CONFIG: "OperationConfig" = {
3435
"static": {
3536
SPANDATA.GEN_AI_SYSTEM: "cohere",
3637
SPANDATA.GEN_AI_OPERATION_NAME: "embeddings",

‎sentry_sdk/integrations/cohere/v1.py‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
if TYPE_CHECKING:
1414
from typing import Any, Callable, Iterator
1515
from cohere import StreamedChatResponse
16+
from sentry_sdk.ai.span_config import OperationConfig
1617

1718
import sentry_sdk
1819
from sentry_sdk.integrations.cohere import (
@@ -37,13 +38,26 @@
3738
_has_chat_types = False
3839

3940

41+
def setup_v1(wrap_embed_fn):
42+
# type: (Callable[..., Any]) -> None
43+
try:
44+
from cohere.base_client import BaseCohere
45+
from cohere.client import Client
46+
except ImportError:
47+
return
48+
49+
BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False)
50+
BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True)
51+
Client.embed = wrap_embed_fn(Client.embed)
52+
53+
4054
def _extract_response_text(response):
4155
# type: (Any) -> list[str] | None
4256
text = getattr(response, "text", None)
4357
return [text] if text is not None else None
4458

4559

46-
COHERE_V1_CHAT_CONFIG = {
60+
COHERE_V1_CHAT_CONFIG: "OperationConfig" = {
4761
"static": {
4862
SPANDATA.GEN_AI_SYSTEM: "cohere",
4963
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
@@ -74,19 +88,6 @@ def _extract_response_text(response):
7488
}
7589

7690

77-
def setup_v1(wrap_embed_fn):
78-
# type: (Callable[..., Any]) -> None
79-
try:
80-
from cohere.base_client import BaseCohere
81-
from cohere.client import Client
82-
except ImportError:
83-
return
84-
85-
BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False)
86-
BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True)
87-
Client.embed = wrap_embed_fn(Client.embed)
88-
89-
9091
def _wrap_chat(f, streaming):
9192
# type: (Callable[..., Any], bool) -> Callable[..., Any]
9293
if not _has_chat_types:
@@ -112,7 +113,7 @@ def new_chat(*args, **kwargs):
112113
origin=CohereIntegration.origin,
113114
) as span:
114115
try:
115-
res = f(*args, **kwargs)
116+
response = f(*args, **kwargs)
116117
except Exception as e:
117118
exc_info = sys.exc_info()
118119
with capture_internal_exceptions():
@@ -129,12 +130,12 @@ def new_chat(*args, **kwargs):
129130
)
130131

131132
if streaming:
132-
return _iter_stream_events(res, span, include_pii)
133+
return _iter_stream_events(response, span, include_pii)
133134
else:
134135
set_response_span_data(
135-
span, res, include_pii, COHERE_V1_CHAT_CONFIG["response"]
136+
span, response, include_pii, COHERE_V1_CHAT_CONFIG["response"]
136137
)
137-
return res
138+
return response
138139

139140
return new_chat
140141

‎sentry_sdk/integrations/cohere/v2.py‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
if TYPE_CHECKING:
1515
from typing import Any, Callable, Iterator
1616
from sentry_sdk.tracing import Span
17+
from sentry_sdk.ai.span_config import OperationConfig
1718

1819
import sentry_sdk
1920
from sentry_sdk.integrations.cohere import (
@@ -44,18 +45,19 @@
4445
except ImportError:
4546
_has_v2 = False
4647

47-
STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")]
4848

49-
_V2_USAGE = {
50-
"input_tokens": [
51-
("usage", "billed_units", "input_tokens"),
52-
("usage", "tokens", "input_tokens"),
53-
],
54-
"output_tokens": [
55-
("usage", "billed_units", "output_tokens"),
56-
("usage", "tokens", "output_tokens"),
57-
],
58-
}
49+
def setup_v2(wrap_embed_fn):
50+
# type: (Callable[..., Any]) -> None
51+
if not _has_v2:
52+
return
53+
CohereV2Client.chat = _wrap_chat_v2(CohereV2Client.chat, streaming=False)
54+
CohereV2Client.chat_stream = _wrap_chat_v2(
55+
CohereV2Client.chat_stream, streaming=True
56+
)
57+
CohereV2Client.embed = wrap_embed_fn(CohereV2Client.embed)
58+
59+
60+
STREAM_DELTA_TEXT_SOURCES = [("delta", "message", "content", "text")]
5961

6062

6163
def _extract_v2_response_text(response):
@@ -68,7 +70,7 @@ def _extract_v2_response_text(response):
6870
return None
6971

7072

71-
COHERE_V2_CHAT_CONFIG = {
73+
COHERE_V2_CHAT_CONFIG: "OperationConfig" = {
7274
"static": {
7375
SPANDATA.GEN_AI_SYSTEM: "cohere",
7476
SPANDATA.GEN_AI_OPERATION_NAME: "chat",
@@ -87,7 +89,16 @@ def _extract_v2_response_text(response):
8789
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: [("message", "tool_calls")],
8890
},
8991
"extract_text": _extract_v2_response_text,
90-
"usage": _V2_USAGE,
92+
"usage": {
93+
"input_tokens": [
94+
("usage", "billed_units", "input_tokens"),
95+
("usage", "tokens", "input_tokens"),
96+
],
97+
"output_tokens": [
98+
("usage", "billed_units", "output_tokens"),
99+
("usage", "tokens", "output_tokens"),
100+
],
101+
},
91102
},
92103
"stream_response": {
93104
"sources": {
@@ -108,17 +119,6 @@ def _extract_v2_response_text(response):
108119
}
109120

110121

111-
def setup_v2(wrap_embed_fn):
112-
# type: (Callable[..., Any]) -> None
113-
if not _has_v2:
114-
return
115-
CohereV2Client.chat = _wrap_chat_v2(CohereV2Client.chat, streaming=False)
116-
CohereV2Client.chat_stream = _wrap_chat_v2(
117-
CohereV2Client.chat_stream, streaming=True
118-
)
119-
CohereV2Client.embed = wrap_embed_fn(CohereV2Client.embed)
120-
121-
122122
def _wrap_chat_v2(f, streaming):
123123
# type: (Callable[..., Any], bool) -> Callable[..., Any]
124124
@wraps(f)
@@ -137,7 +137,7 @@ def new_chat(*args, **kwargs):
137137
origin=CohereIntegration.origin,
138138
) as span:
139139
try:
140-
res = f(*args, **kwargs)
140+
response = f(*args, **kwargs)
141141
except Exception as e:
142142
exc_info = sys.exc_info()
143143
with capture_internal_exceptions():
@@ -153,11 +153,11 @@ def new_chat(*args, **kwargs):
153153
span, kwargs, integration, COHERE_V2_CHAT_CONFIG, span_data
154154
)
155155
if streaming:
156-
return _iter_v2_stream_events(res, span, include_pii)
156+
return _iter_v2_stream_events(response, span, include_pii)
157157
set_response_span_data(
158-
span, res, include_pii, COHERE_V2_CHAT_CONFIG["response"]
158+
span, response, include_pii, COHERE_V2_CHAT_CONFIG["response"]
159159
)
160-
return res
160+
return response
161161

162162
return new_chat
163163

0 commit comments

Comments
 (0)