diff --git a/sentry_sdk/integrations/typesafe.py b/sentry_sdk/integrations/typesafe.py index e0526cde62..ab6586dc7d 100644 --- a/sentry_sdk/integrations/typesafe.py +++ b/sentry_sdk/integrations/typesafe.py @@ -47,12 +47,29 @@ class InputMessageModel(TypedDict): state: NotRequired[JSONContent] questions: NotRequired[dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]] + class NoulEvaluationModel(TypedDict): + type: Literal["noul"] + noul: float + + class ChoiceEvaluationModel(TypedDict): + type: Literal["choice"] + choice: str + probabilities: dict[str, float] + confidence: float + + class ScoreEvaluationModel(TypedDict): + type: Literal["score"] + score: float + probabilities: dict[int, float] + confidence: float + legend: dict[int, Union[str, dict[str, Any], list[Any]]] + try: - from typesafe_sdk import Choice, Noul, Score + from typesafe_sdk import Choice, ChoiceAnswer, Noul, NoulAnswer, Score, ScoreAnswer from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient from typesafe_sdk._core.client.sync.client import TypeSafeClient - from typesafe_sdk._core.response_types import SystemOneResponse + from typesafe_sdk._core.response_types import Answer, SystemOneResponse except ImportError: raise DidNotEnable("typesafe-sdk not installed") @@ -152,6 +169,42 @@ def _transform_questions( return transformed_questions +def _transform_evaluation_answers( + answers: "dict[str, Answer]", +) -> ( + "dict[str, Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]]" +): + items: "dict[str, Union[NoulEvaluationModel, ChoiceEvaluationModel, ScoreEvaluationModel]]" = {} + for name, answer in answers.items(): + if isinstance(answer, NoulAnswer): + items[name] = { + "type": "noul", + "noul": answer.noul, + } + continue + + if isinstance(answer, ChoiceAnswer): + items[name] = { + "type": "choice", + "choice": answer.choice, + "probabilities": answer.probabilities, + "confidence": answer.confidence, + } + continue + + if isinstance(answer, ScoreAnswer): + items[name] = { + "type": "score", + "score": answer.score, + "probabilities": answer.probabilities, + "confidence": answer.confidence, + "legend": answer.legend, + } + continue + + return items + + def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) 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 SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens ) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["outputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + set_on_span( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps( + [ + { + "type": "evaluation", + "answers": _transform_evaluation_answers( + response.answers + ), + } + ] + ), + ) + return response return wrap_system_one @@ -317,6 +391,27 @@ async def wrap_system_one_async( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens ) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["outputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + set_on_span( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps( + [ + { + "type": "evaluation", + "answers": _transform_evaluation_answers( + response.answers + ), + } + ] + ), + ) + return response return wrap_system_one_async diff --git a/tests/integrations/typesafe/test_typesafe.py b/tests/integrations/typesafe/test_typesafe.py index 24948e4afe..bb0f7bd06b 100644 --- a/tests/integrations/typesafe/test_typesafe.py +++ b/tests/integrations/typesafe/test_typesafe.py @@ -164,6 +164,48 @@ def test_system_one( }, } ] + + assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "type": "evaluation", + "answers": { + "spam": { + "type": "noul", + "noul": 0.98, + }, + "tone": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + "spam_obj": { + "type": "noul", + "noul": 0.98, + }, + "tone_obj": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency_obj": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + }, + } + ] else: items = capture_items("transaction") @@ -251,6 +293,48 @@ def test_system_one( } ] + assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "type": "evaluation", + "answers": { + "spam": { + "type": "noul", + "noul": 0.98, + }, + "tone": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + "spam_obj": { + "type": "noul", + "noul": 0.98, + }, + "tone_obj": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency_obj": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + }, + } + ] + @pytest.mark.asyncio @pytest.mark.parametrize("span_streaming", [True, False]) @@ -366,6 +450,48 @@ async def test_system_one_async( }, } ] + + assert json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "type": "evaluation", + "answers": { + "spam": { + "type": "noul", + "noul": 0.98, + }, + "tone": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + "spam_obj": { + "type": "noul", + "noul": 0.98, + }, + "tone_obj": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency_obj": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + }, + } + ] else: items = capture_items("transaction") @@ -455,3 +581,45 @@ async def test_system_one_async( }, } ] + + assert json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "type": "evaluation", + "answers": { + "spam": { + "type": "noul", + "noul": 0.98, + }, + "tone": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + "spam_obj": { + "type": "noul", + "noul": 0.98, + }, + "tone_obj": { + "type": "choice", + "choice": "friendly", + "probabilities": {"friendly": 0.9, "hostile": 0.1}, + "confidence": 0.9, + }, + "urgency_obj": { + "type": "score", + "score": 1.7, + "probabilities": {"0": 0.1, "1": 0.1, "2": 0.8}, + "confidence": 0.8, + "legend": {"0": "can wait", "1": "this week", "2": "today"}, + }, + }, + } + ]