Skip to content

Commit dfc004a

Browse files
alexander-alderman-webbalexander-alderman-webb
authored
feat(typesafe): Record gen_ai.output.messages (#7720)
Handle `NoulAnswer`, `ChoiceAnswer` and `ScoreAnswer` Pydantic objects.
1 parent d36887b commit dfc004a

2 files changed

Lines changed: 265 additions & 2 deletions

File tree

‎sentry_sdk/integrations/typesafe.py‎

Lines changed: 97 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,29 @@ class InputMessageModel(TypedDict):
4747
state: NotRequired[JSONContent]
4848
questions: NotRequired[dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]]
4949

50+
class NoulEvaluationModel(TypedDict):
51+
type: Literal["noul"]
52+
noul: float
53+
54+
class ChoiceEvaluationModel(TypedDict):
55+
type: Literal["choice"]
56+
choice: str
57+
probabilities: dict[str, float]
58+
confidence: float
59+
60+
class ScoreEvaluationModel(TypedDict):
61+
type: Literal["score"]
62+
score: float
63+
probabilities: dict[int, float]
64+
confidence: float
65+
legend: dict[int, Union[str, dict[str, Any], list[Any]]]
66+
5067

5168
try:
52-
from typesafe_sdk import Choice, Noul, Score
69+
from typesafe_sdk import Choice, ChoiceAnswer, Noul, NoulAnswer, Score, ScoreAnswer
5370
from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient
5471
from typesafe_sdk._core.client.sync.client import TypeSafeClient
55-
from typesafe_sdk._core.response_types import SystemOneResponse
72+
from typesafe_sdk._core.response_types import Answer, SystemOneResponse
5673
except ImportError:
5774
raise DidNotEnable("typesafe-sdk not installed")
5875

@@ -152,6 +169,42 @@ def _transform_questions(
152169
return transformed_questions
153170

154171

172+
def _transform_evaluation_answers(
173+
answers: "dict[str, Answer]",
174+
) -> (
175+
"dict[str, Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]]"
176+
):
177+
items: "dict[str, Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]]" = {}
178+
for name, answer in answers.items():
179+
if isinstance(answer, NoulAnswer):
180+
items[name] = {
181+
"type": "noul",
182+
"noul": answer.noul,
183+
}
184+
continue
185+
186+
if isinstance(answer, ChoiceAnswer):
187+
items[name] = {
188+
"type": "choice",
189+
"choice": answer.choice,
190+
"probabilities": answer.probabilities,
191+
"confidence": answer.confidence,
192+
}
193+
continue
194+
195+
if isinstance(answer, ScoreAnswer):
196+
items[name] = {
197+
"type": "score",
198+
"score": answer.score,
199+
"probabilities": answer.probabilities,
200+
"confidence": answer.confidence,
201+
"legend": answer.legend,
202+
}
203+
continue
204+
205+
return items
206+
207+
155208
def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]":
156209
@wraps(f)
157210
def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "Any":
@@ -231,6 +284,27 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
231284
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
232285
)
233286

287+
if (
288+
has_data_collection_enabled(client.options)
289+
and client.options["data_collection"]["gen_ai"]["outputs"]
290+
) or (
291+
not has_data_collection_enabled(client.options)
292+
and should_send_default_pii()
293+
):
294+
set_on_span(
295+
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
296+
json.dumps(
297+
[
298+
{
299+
"type": "evaluation",
300+
"answers": _transform_evaluation_answers(
301+
response.answers
302+
),
303+
}
304+
]
305+
),
306+
)
307+
234308
return response
235309

236310
return wrap_system_one
@@ -317,6 +391,27 @@ async def wrap_system_one_async(
317391
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
318392
)
319393

394+
if (
395+
has_data_collection_enabled(client.options)
396+
and client.options["data_collection"]["gen_ai"]["outputs"]
397+
) or (
398+
not has_data_collection_enabled(client.options)
399+
and should_send_default_pii()
400+
):
401+
set_on_span(
402+
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
403+
json.dumps(
404+
[
405+
{
406+
"type": "evaluation",
407+
"answers": _transform_evaluation_answers(
408+
response.answers
409+
),
410+
}
411+
]
412+
),
413+
)
414+
320415
return response
321416

322417
return wrap_system_one_async

‎tests/integrations/typesafe/test_typesafe.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,48 @@ def test_system_one(
164164
},
165165
}
166166
]
167+
168+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
169+
{
170+
"type": "evaluation",
171+
"answers": {
172+
"spam": {
173+
"type": "noul",
174+
"noul": 0.98,
175+
},
176+
"tone": {
177+
"type": "choice",
178+
"choice": "friendly",
179+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
180+
"confidence": 0.9,
181+
},
182+
"urgency": {
183+
"type": "score",
184+
"score": 1.7,
185+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
186+
"confidence": 0.8,
187+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
188+
},
189+
"spam_obj": {
190+
"type": "noul",
191+
"noul": 0.98,
192+
},
193+
"tone_obj": {
194+
"type": "choice",
195+
"choice": "friendly",
196+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
197+
"confidence": 0.9,
198+
},
199+
"urgency_obj": {
200+
"type": "score",
201+
"score": 1.7,
202+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
203+
"confidence": 0.8,
204+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
205+
},
206+
},
207+
}
208+
]
167209
else:
168210
items = capture_items("transaction")
169211

@@ -251,6 +293,48 @@ def test_system_one(
251293
}
252294
]
253295

296+
assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
297+
{
298+
"type": "evaluation",
299+
"answers": {
300+
"spam": {
301+
"type": "noul",
302+
"noul": 0.98,
303+
},
304+
"tone": {
305+
"type": "choice",
306+
"choice": "friendly",
307+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
308+
"confidence": 0.9,
309+
},
310+
"urgency": {
311+
"type": "score",
312+
"score": 1.7,
313+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
314+
"confidence": 0.8,
315+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
316+
},
317+
"spam_obj": {
318+
"type": "noul",
319+
"noul": 0.98,
320+
},
321+
"tone_obj": {
322+
"type": "choice",
323+
"choice": "friendly",
324+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
325+
"confidence": 0.9,
326+
},
327+
"urgency_obj": {
328+
"type": "score",
329+
"score": 1.7,
330+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
331+
"confidence": 0.8,
332+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
333+
},
334+
},
335+
}
336+
]
337+
254338

255339
@pytest.mark.asyncio
256340
@pytest.mark.parametrize("span_streaming", [True, False])
@@ -366,6 +450,48 @@ async def test_system_one_async(
366450
},
367451
}
368452
]
453+
454+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
455+
{
456+
"type": "evaluation",
457+
"answers": {
458+
"spam": {
459+
"type": "noul",
460+
"noul": 0.98,
461+
},
462+
"tone": {
463+
"type": "choice",
464+
"choice": "friendly",
465+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
466+
"confidence": 0.9,
467+
},
468+
"urgency": {
469+
"type": "score",
470+
"score": 1.7,
471+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
472+
"confidence": 0.8,
473+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
474+
},
475+
"spam_obj": {
476+
"type": "noul",
477+
"noul": 0.98,
478+
},
479+
"tone_obj": {
480+
"type": "choice",
481+
"choice": "friendly",
482+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
483+
"confidence": 0.9,
484+
},
485+
"urgency_obj": {
486+
"type": "score",
487+
"score": 1.7,
488+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
489+
"confidence": 0.8,
490+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
491+
},
492+
},
493+
}
494+
]
369495
else:
370496
items = capture_items("transaction")
371497

@@ -455,3 +581,45 @@ async def test_system_one_async(
455581
},
456582
}
457583
]
584+
585+
assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
586+
{
587+
"type": "evaluation",
588+
"answers": {
589+
"spam": {
590+
"type": "noul",
591+
"noul": 0.98,
592+
},
593+
"tone": {
594+
"type": "choice",
595+
"choice": "friendly",
596+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
597+
"confidence": 0.9,
598+
},
599+
"urgency": {
600+
"type": "score",
601+
"score": 1.7,
602+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
603+
"confidence": 0.8,
604+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
605+
},
606+
"spam_obj": {
607+
"type": "noul",
608+
"noul": 0.98,
609+
},
610+
"tone_obj": {
611+
"type": "choice",
612+
"choice": "friendly",
613+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
614+
"confidence": 0.9,
615+
},
616+
"urgency_obj": {
617+
"type": "score",
618+
"score": 1.7,
619+
"probabilities": {"0": 0.1, "1": 0.1, "2": 0.8},
620+
"confidence": 0.8,
621+
"legend": {"0": "can wait", "1": "this week", "2": "today"},
622+
},
623+
},
624+
}
625+
]

0 commit comments

Comments
 (0)