Skip to content

Commit 16ef9d3

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
feat(typesafe): Record gen_ai.output.messages
1 parent 17261c5 commit 16ef9d3

2 files changed

Lines changed: 247 additions & 2 deletions

File tree

‎sentry_sdk/integrations/typesafe.py‎

Lines changed: 79 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,32 @@ class ScoreModel(TypedDict):
4040
criteria: NotRequired[Sequence[JSONContent]]
4141
instructions: NotRequired[JSONContent | None]
4242

43+
class NoulEvaluationModel(TypedDict):
44+
name: str
45+
type: Literal["noul"]
46+
noul: float
47+
48+
class ChoiceEvaluationModel(TypedDict):
49+
name: str
50+
type: Literal["choice"]
51+
choice: str
52+
probabilities: dict[str, float]
53+
confidence: float
54+
55+
class ScoreEvaluationModel(TypedDict):
56+
name: str
57+
type: Literal["score"]
58+
score: float
59+
probabilities: list[float]
60+
confidence: float
61+
legend: list[JSONContent]
62+
4363

4464
try:
45-
from typesafe_sdk import Choice, Noul, Score
65+
from typesafe_sdk import Choice, ChoiceAnswer, Noul, NoulAnswer, Score, ScoreAnswer
4666
from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient
4767
from typesafe_sdk._core.client.sync.client import TypeSafeClient
48-
from typesafe_sdk._core.response_types import SystemOneResponse
68+
from typesafe_sdk._core.response_types import Answer, SystemOneResponse
4969
except ImportError:
5070
raise DidNotEnable("typesafe-sdk not installed")
5171

@@ -151,6 +171,53 @@ def _transform_questions(
151171
return transformed_questions
152172

153173

174+
def _transform_evaluation_answers(
175+
answers: dict[str, Answer],
176+
) -> list[Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]]:
177+
items: list[
178+
Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]
179+
] = []
180+
for name, answer in answers.items():
181+
if isinstance(answer, NoulAnswer):
182+
items.append(
183+
{
184+
"name": name,
185+
"type": "noul",
186+
"noul": answer.noul,
187+
}
188+
)
189+
continue
190+
191+
if isinstance(answer, ChoiceAnswer):
192+
items.append(
193+
{
194+
"name": name,
195+
"type": "choice",
196+
"choice": answer.choice,
197+
"probabilities": answer.probabilities,
198+
"confidence": answer.confidence,
199+
}
200+
)
201+
continue
202+
203+
if isinstance(answer, ScoreAnswer):
204+
items.append(
205+
{
206+
"name": name,
207+
"type": "score",
208+
"score": answer.score,
209+
"probabilities": [
210+
value for _, value in answer.probabilities.items()
211+
],
212+
"confidence": answer.confidence,
213+
"legend": [value for _, value in answer.legend.items()],
214+
}
215+
)
216+
continue
217+
218+
return items
219+
220+
154221
def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]":
155222
@wraps(f)
156223
def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "Any":
@@ -222,6 +289,11 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
222289
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
223290
)
224291

292+
set_on_span(
293+
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
294+
json.dumps(_transform_evaluation_answers(response.answers)),
295+
)
296+
225297
return response
226298

227299
return wrap_system_one
@@ -300,6 +372,11 @@ async def wrap_system_one_async(
300372
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
301373
)
302374

375+
set_on_span(
376+
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
377+
json.dumps(_transform_evaluation_answers(response.answers)),
378+
)
379+
303380
return response
304381

305382
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
@@ -167,6 +167,48 @@ def test_system_one(
167167
},
168168
],
169169
}
170+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
171+
{
172+
"type": "noul",
173+
"name": "spam",
174+
"noul": 0.98,
175+
},
176+
{
177+
"type": "choice",
178+
"name": "tone",
179+
"choice": "friendly",
180+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
181+
"confidence": 0.9,
182+
},
183+
{
184+
"type": "score",
185+
"name": "urgency",
186+
"score": 1.7,
187+
"probabilities": [0.1, 0.1, 0.8],
188+
"confidence": 0.8,
189+
"legend": ["can wait", "this week", "today"],
190+
},
191+
{
192+
"type": "noul",
193+
"name": "spam_obj",
194+
"noul": 0.98,
195+
},
196+
{
197+
"type": "choice",
198+
"name": "tone_obj",
199+
"choice": "friendly",
200+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
201+
"confidence": 0.9,
202+
},
203+
{
204+
"type": "score",
205+
"name": "urgency_obj",
206+
"score": 1.7,
207+
"probabilities": [0.1, 0.1, 0.8],
208+
"confidence": 0.8,
209+
"legend": ["can wait", "this week", "today"],
210+
},
211+
]
170212
else:
171213
items = capture_items("transaction")
172214

@@ -257,6 +299,48 @@ def test_system_one(
257299
},
258300
],
259301
}
302+
assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
303+
{
304+
"type": "noul",
305+
"name": "spam",
306+
"noul": 0.98,
307+
},
308+
{
309+
"type": "choice",
310+
"name": "tone",
311+
"choice": "friendly",
312+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
313+
"confidence": 0.9,
314+
},
315+
{
316+
"type": "score",
317+
"name": "urgency",
318+
"score": 1.7,
319+
"probabilities": [0.1, 0.1, 0.8],
320+
"confidence": 0.8,
321+
"legend": ["can wait", "this week", "today"],
322+
},
323+
{
324+
"type": "noul",
325+
"name": "spam_obj",
326+
"noul": 0.98,
327+
},
328+
{
329+
"type": "choice",
330+
"name": "tone_obj",
331+
"choice": "friendly",
332+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
333+
"confidence": 0.9,
334+
},
335+
{
336+
"type": "score",
337+
"name": "urgency_obj",
338+
"score": 1.7,
339+
"probabilities": [0.1, 0.1, 0.8],
340+
"confidence": 0.8,
341+
"legend": ["can wait", "this week", "today"],
342+
},
343+
]
260344

261345

262346
@pytest.mark.asyncio
@@ -376,6 +460,48 @@ async def test_system_one_async(
376460
},
377461
],
378462
}
463+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
464+
{
465+
"type": "noul",
466+
"name": "spam",
467+
"noul": 0.98,
468+
},
469+
{
470+
"type": "choice",
471+
"name": "tone",
472+
"choice": "friendly",
473+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
474+
"confidence": 0.9,
475+
},
476+
{
477+
"type": "score",
478+
"name": "urgency",
479+
"score": 1.7,
480+
"probabilities": [0.1, 0.1, 0.8],
481+
"confidence": 0.8,
482+
"legend": ["can wait", "this week", "today"],
483+
},
484+
{
485+
"type": "noul",
486+
"name": "spam_obj",
487+
"noul": 0.98,
488+
},
489+
{
490+
"type": "choice",
491+
"name": "tone_obj",
492+
"choice": "friendly",
493+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
494+
"confidence": 0.9,
495+
},
496+
{
497+
"type": "score",
498+
"name": "urgency_obj",
499+
"score": 1.7,
500+
"probabilities": [0.1, 0.1, 0.8],
501+
"confidence": 0.8,
502+
"legend": ["can wait", "this week", "today"],
503+
},
504+
]
379505
else:
380506
items = capture_items("transaction")
381507

@@ -469,3 +595,45 @@ async def test_system_one_async(
469595
},
470596
],
471597
}
598+
assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [
599+
{
600+
"type": "noul",
601+
"name": "spam",
602+
"noul": 0.98,
603+
},
604+
{
605+
"type": "choice",
606+
"name": "tone",
607+
"choice": "friendly",
608+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
609+
"confidence": 0.9,
610+
},
611+
{
612+
"type": "score",
613+
"name": "urgency",
614+
"score": 1.7,
615+
"probabilities": [0.1, 0.1, 0.8],
616+
"confidence": 0.8,
617+
"legend": ["can wait", "this week", "today"],
618+
},
619+
{
620+
"type": "noul",
621+
"name": "spam_obj",
622+
"noul": 0.98,
623+
},
624+
{
625+
"type": "choice",
626+
"name": "tone_obj",
627+
"choice": "friendly",
628+
"probabilities": {"friendly": 0.9, "hostile": 0.1},
629+
"confidence": 0.9,
630+
},
631+
{
632+
"type": "score",
633+
"name": "urgency_obj",
634+
"score": 1.7,
635+
"probabilities": [0.1, 0.1, 0.8],
636+
"confidence": 0.8,
637+
"legend": ["can wait", "this week", "today"],
638+
},
639+
]

0 commit comments

Comments
 (0)