diff --git a/sentry_sdk/integrations/typesafe.py b/sentry_sdk/integrations/typesafe.py index 810dcfa057..796cd33fca 100644 --- a/sentry_sdk/integrations/typesafe.py +++ b/sentry_sdk/integrations/typesafe.py @@ -1,16 +1,55 @@ +import json +from collections.abc import Mapping from functools import wraps -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import sentry_sdk from sentry_sdk.ai.utils import get_start_span_function from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.tracing_utils import has_span_streaming_enabled +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import ( + has_data_collection_enabled, + has_span_streaming_enabled, +) if TYPE_CHECKING: - from typing import Any, Callable + from typing import ( + Any, + Callable, + Literal, + Mapping, + NotRequired, + Sequence, + TypedDict, + Union, + ) + + from typesafe_sdk import JSONContent + from typesafe_sdk._core.question_types import Questions + + class NoulModel(TypedDict): + type: Literal["noul"] + instructions: NotRequired[JSONContent | None] + + class ChoiceModel(TypedDict): + type: Literal["choice"] + criteria: NotRequired[Mapping[str, JSONContent | None]] + instructions: NotRequired[JSONContent | None] + + class ScoreModel(TypedDict): + type: Literal["score"] + criteria: NotRequired[Sequence[JSONContent]] + instructions: NotRequired[JSONContent | None] + + class InputMessageModel(TypedDict): + type: Literal["evaluation"] + state: NotRequired[JSONContent] + questions: NotRequired[dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]] + try: + from typesafe_sdk import Choice, Noul, Score 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 @@ -30,6 +69,89 @@ def setup_once() -> None: ) +def _transform_questions( + questions: "Questions", +) -> "dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]": + transformed_questions: "dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]" = {} + for name, question in questions.items(): + if isinstance(question, Noul): + noul: "NoulModel" = { + "type": "noul", + } + + if question.instructions is not None: + noul["instructions"] = question.instructions + + transformed_questions[name] = noul + continue + + if isinstance(question, Choice): + choice: "ChoiceModel" = { + "type": "choice", + "criteria": question.criteria, + } + + if question.instructions is not None: + choice["instructions"] = question.instructions + + transformed_questions[name] = choice + continue + + if isinstance(question, Score): + score: "ScoreModel" = { + "type": "score", + "criteria": question.criteria, + } + + if question.instructions is not None: + score["instructions"] = question.instructions + + transformed_questions[name] = score + continue + + if not isinstance(question, dict): + continue + + question_type = question.get("type") + if question_type == "noul": + noul = { + "type": question_type, + } + if "instructions" in question: + noul["instructions"] = question["instructions"] + + transformed_questions[name] = noul + continue + + if question_type == "choice": + choice = { + "type": question_type, + } + if "criteria" in question: + choice["criteria"] = cast( + "Mapping[str, JSONContent | None]", question["criteria"] + ) + if "instructions" in question: + choice["instructions"] = question["instructions"] + + transformed_questions[name] = choice + continue + + if question_type == "score": + score = { + "type": question_type, + } + if "criteria" in question: + score["criteria"] = cast("Sequence[JSONContent]", question["criteria"]) + if "instructions" in question: + score["instructions"] = question["instructions"] + + transformed_questions[name] = score + continue + + return transformed_questions + + def _wrap_system_one(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "Any": @@ -71,6 +193,27 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A if model is not None: set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["inputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + input_message: "InputMessageModel" = { + "type": "evaluation", + } + + state = args[0] if len(args) > 0 else kwargs.get("state") + if state is not None: + input_message["state"] = state + + questions = args[1] if len(args) > 1 else kwargs.get("questions") + if isinstance(questions, Mapping): + input_message["questions"] = _transform_questions(questions) + + set_on_span(SPANDATA.GEN_AI_INPUT_MESSAGES, json.dumps([input_message])) + response = f(self, *args, **kwargs) if not isinstance(response, SystemOneResponse): @@ -126,6 +269,27 @@ async def wrap_system_one_async( if model is not None: set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["inputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + input_message: "InputMessageModel" = { + "type": "evaluation", + } + + state = args[0] if len(args) > 0 else kwargs.get("state") + if state is not None: + input_message["state"] = state + + questions = args[1] if len(args) > 1 else kwargs.get("questions") + if isinstance(questions, Mapping): + input_message["questions"] = _transform_questions(questions) + + set_on_span(SPANDATA.GEN_AI_INPUT_MESSAGES, json.dumps([input_message])) + response = await f(self, *args, **kwargs) if not isinstance(response, SystemOneResponse): diff --git a/tests/integrations/typesafe/test_typesafe.py b/tests/integrations/typesafe/test_typesafe.py index 4c652c22ae..08fad17e29 100644 --- a/tests/integrations/typesafe/test_typesafe.py +++ b/tests/integrations/typesafe/test_typesafe.py @@ -1,3 +1,4 @@ +import json from unittest import mock import pytest @@ -65,6 +66,7 @@ def test_system_one( sentry_init( integrations=[TypeSafeIntegration()], traces_sample_rate=1.0, + send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -119,6 +121,46 @@ def test_system_one( assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" + + assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [ + { + "type": "evaluation", + "state": { + "subject": "Charged twice this month", + "body": "I see two charges of $49. I only have one account. Please fix this ASAP.", + }, + "questions": { + "spam": { + "type": "noul", + "instructions": "Spam?", + }, + "tone": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + "spam_obj": { + "type": "noul", + "instructions": "Spam?", + }, + "tone_obj": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality_obj": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + }, + } + ] else: items = capture_items("transaction") @@ -163,6 +205,46 @@ def test_system_one( assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" + assert json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [ + { + "type": "evaluation", + "state": { + "subject": "Charged twice this month", + "body": "I see two charges of $49. I only have one account. Please fix this ASAP.", + }, + "questions": { + "spam": { + "type": "noul", + "instructions": "Spam?", + }, + "tone": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + "spam_obj": { + "type": "noul", + "instructions": "Spam?", + }, + "tone_obj": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality_obj": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + }, + } + ] + @pytest.mark.asyncio @pytest.mark.parametrize("span_streaming", [True, False]) @@ -177,6 +259,7 @@ async def test_system_one_async( sentry_init( integrations=[TypeSafeIntegration()], traces_sample_rate=1.0, + send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -234,6 +317,46 @@ async def test_system_one_async( assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" + + assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [ + { + "type": "evaluation", + "state": { + "subject": "Charged twice this month", + "body": "I see two charges of $49. I only have one account. Please fix this ASAP.", + }, + "questions": { + "spam": { + "type": "noul", + "instructions": "Spam?", + }, + "tone": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + "spam_obj": { + "type": "noul", + "instructions": "Spam?", + }, + "tone_obj": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality_obj": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + }, + } + ] else: items = capture_items("transaction") @@ -280,3 +403,43 @@ async def test_system_one_async( assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" + + assert json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [ + { + "type": "evaluation", + "state": { + "subject": "Charged twice this month", + "body": "I see two charges of $49. I only have one account. Please fix this ASAP.", + }, + "questions": { + "spam": { + "type": "noul", + "instructions": "Spam?", + }, + "tone": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + "spam_obj": { + "type": "noul", + "instructions": "Spam?", + }, + "tone_obj": { + "type": "choice", + "instructions": "Tone?", + "criteria": {"friendly": None, "hostile": None}, + }, + "quality_obj": { + "type": "score", + "instructions": "Quality?", + "criteria": ["bad", "ok", "great"], + }, + }, + } + ]