diff --git a/sentry_sdk/integrations/typesafe.py b/sentry_sdk/integrations/typesafe.py index e81ce788a9..810dcfa057 100644 --- a/sentry_sdk/integrations/typesafe.py +++ b/sentry_sdk/integrations/typesafe.py @@ -13,6 +13,7 @@ try: 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 except ImportError: raise DidNotEnable("typesafe-sdk not installed") @@ -70,7 +71,14 @@ 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) - return f(self, *args, **kwargs) + response = f(self, *args, **kwargs) + + if not isinstance(response, SystemOneResponse): + return response + + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model) + + return response return wrap_system_one @@ -118,6 +126,13 @@ async def wrap_system_one_async( if model is not None: set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) - return await f(self, *args, **kwargs) + response = await f(self, *args, **kwargs) + + if not isinstance(response, SystemOneResponse): + return response + + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model) + + return response return wrap_system_one_async diff --git a/tests/integrations/typesafe/test_typesafe.py b/tests/integrations/typesafe/test_typesafe.py index c869fd7a3a..4c652c22ae 100644 --- a/tests/integrations/typesafe/test_typesafe.py +++ b/tests/integrations/typesafe/test_typesafe.py @@ -117,6 +117,8 @@ def test_system_one( assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "typesafe" assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "evaluate" assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" + + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" else: items = capture_items("transaction") @@ -159,6 +161,8 @@ def test_system_one( assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "evaluate" assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" + @pytest.mark.asyncio @pytest.mark.parametrize("span_streaming", [True, False]) @@ -228,6 +232,8 @@ async def test_system_one_async( assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "typesafe" assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "evaluate" assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" + + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest" else: items = capture_items("transaction") @@ -272,3 +278,5 @@ async def test_system_one_async( assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "typesafe" assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "evaluate" assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "jev-latest" + + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest"