Skip to content

Commit 76e8c03

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
feat(typesafe): Handle inline images
1 parent 7ad4e1a commit 76e8c03

4 files changed

Lines changed: 2432 additions & 2253 deletions

File tree

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ typing = [
8181
"aiohttp>=3.14.3",
8282
"openai-agents>=0.22.0",
8383
"mistralai>=2.10.1",
84-
"typesafe-sdk>=0.7.2",
84+
"typesafe-sdk>=0.7.5",
8585
]
8686
test = [
8787
"dataclasses ; python_full_version < '3.7'",

‎sentry_sdk/integrations/typesafe.py‎

Lines changed: 51 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,26 +25,26 @@
2525
Union,
2626
)
2727

28-
from typesafe_sdk import JSONContent
28+
from typesafe_sdk import Content
2929
from typesafe_sdk._core.question_types import Questions
3030

3131
class NoulModel(TypedDict):
3232
type: Literal["noul"]
33-
instructions: NotRequired[JSONContent | None]
33+
instructions: NotRequired[Content | None]
3434

3535
class ChoiceModel(TypedDict):
3636
type: Literal["choice"]
37-
criteria: NotRequired[Mapping[str, JSONContent | None]]
38-
instructions: NotRequired[JSONContent | None]
37+
criteria: NotRequired[Mapping[str, Content | None]]
38+
instructions: NotRequired[Content | None]
3939

4040
class ScoreModel(TypedDict):
4141
type: Literal["score"]
42-
criteria: NotRequired[Sequence[JSONContent]]
43-
instructions: NotRequired[JSONContent | None]
42+
criteria: NotRequired[Sequence[Content]]
43+
instructions: NotRequired[Content | None]
4444

4545
class InputMessageModel(TypedDict):
4646
type: Literal["evaluation"]
47-
state: NotRequired[JSONContent]
47+
state: NotRequired[Content]
4848
questions: NotRequired[dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]]
4949

5050
class NoulEvaluationModel(TypedDict):
@@ -66,14 +66,27 @@ class ScoreEvaluationModel(TypedDict):
6666

6767

6868
try:
69-
from typesafe_sdk import Choice, ChoiceAnswer, Noul, NoulAnswer, Score, ScoreAnswer
69+
from typesafe_sdk import (
70+
Choice,
71+
ChoiceAnswer,
72+
Noul,
73+
NoulAnswer,
74+
Score,
75+
ScoreAnswer,
76+
)
7077
from typesafe_sdk._core.client.aio.client import AsyncTypeSafeClient
7178
from typesafe_sdk._core.client.sync.client import TypeSafeClient
7279
from typesafe_sdk._core.response_types import Answer, SystemOneResponse
7380
except ImportError:
7481
raise DidNotEnable("typesafe-sdk not installed")
7582

7683

84+
try:
85+
from typesafe_sdk import Image
86+
except ImportError:
87+
Image = None # type: ignore[misc, assignment]
88+
89+
7790
class TypeSafeIntegration(Integration):
7891
identifier = "typesafe"
7992
origin = f"auto.ai.{identifier}"
@@ -86,6 +99,16 @@ def setup_once() -> None:
8699
)
87100

88101

102+
def _json_default(value: object) -> str:
103+
"""
104+
Render an inline typesafe `Image` as its reference, mirroring the SDK's handling.
105+
"""
106+
if Image is not None and isinstance(value, Image):
107+
return str(value)
108+
109+
return "<unserializable>"
110+
111+
89112
def _transform_questions(
90113
questions: "Questions",
91114
) -> "dict[str, Union[NoulModel, ChoiceModel, ScoreModel]]":
@@ -146,7 +169,7 @@ def _transform_questions(
146169
}
147170
if "criteria" in question:
148171
choice["criteria"] = cast(
149-
"Mapping[str, JSONContent | None]", question["criteria"]
172+
"Mapping[str, Content | None]", question["criteria"]
150173
)
151174
if "instructions" in question:
152175
choice["instructions"] = question["instructions"]
@@ -159,7 +182,7 @@ def _transform_questions(
159182
"type": question_type,
160183
}
161184
if "criteria" in question:
162-
score["criteria"] = cast("Sequence[JSONContent]", question["criteria"])
185+
score["criteria"] = cast("Sequence[Content]", question["criteria"])
163186
if "instructions" in question:
164187
score["instructions"] = question["instructions"]
165188

@@ -265,7 +288,13 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
265288
if isinstance(questions, Mapping):
266289
input_message["questions"] = _transform_questions(questions)
267290

268-
set_on_span(SPANDATA.GEN_AI_INPUT_MESSAGES, json.dumps([input_message]))
291+
set_on_span(
292+
SPANDATA.GEN_AI_INPUT_MESSAGES,
293+
json.dumps(
294+
[input_message],
295+
default=_json_default,
296+
),
297+
)
269298

270299
response = f(self, *args, **kwargs)
271300

@@ -301,7 +330,8 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
301330
response.answers
302331
),
303332
}
304-
]
333+
],
334+
default=_json_default,
305335
),
306336
)
307337

@@ -372,7 +402,13 @@ async def wrap_system_one_async(
372402
if isinstance(questions, Mapping):
373403
input_message["questions"] = _transform_questions(questions)
374404

375-
set_on_span(SPANDATA.GEN_AI_INPUT_MESSAGES, json.dumps([input_message]))
405+
set_on_span(
406+
SPANDATA.GEN_AI_INPUT_MESSAGES,
407+
json.dumps(
408+
[input_message],
409+
default=_json_default,
410+
),
411+
)
376412

377413
response = await f(self, *args, **kwargs)
378414

@@ -408,7 +444,8 @@ async def wrap_system_one_async(
408444
response.answers
409445
),
410446
}
411-
]
447+
],
448+
default=_json_default,
412449
),
413450
)
414451

‎tests/integrations/typesafe/test_typesafe.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@
1515
Usage,
1616
)
1717

18+
try:
19+
from typesafe_sdk import Image
20+
except ImportError:
21+
Image = None
22+
1823
import sentry_sdk
1924
from sentry_sdk.consts import OP, SPANDATA
2025
from sentry_sdk.integrations.typesafe import TypeSafeIntegration
@@ -623,3 +628,137 @@ async def test_system_one_async(
623628
},
624629
}
625630
]
631+
632+
633+
@pytest.mark.skipif(
634+
Image is None, reason="Image was added in version v0.7.5 of typesafe-sdk"
635+
)
636+
def test_system_one_with_inline_image(
637+
sentry_init,
638+
capture_items,
639+
typesafe_response,
640+
):
641+
sentry_init(
642+
integrations=[TypeSafeIntegration()],
643+
traces_sample_rate=1.0,
644+
send_default_pii=True,
645+
trace_lifecycle="stream",
646+
)
647+
648+
photo = Image(b"\x89PNG\r\n\x1a\nimage data")
649+
listing = Image(b"\xff\xd8\xffimage data")
650+
651+
client = TypeSafeClient(api_key="z")
652+
items = capture_items("span")
653+
654+
with mock.patch.object(
655+
TypeSafeClient, "_request", return_value=typesafe_response
656+
), sentry_sdk.start_transaction(name="typesafe"):
657+
client.system_one(
658+
state=[
659+
{"role": "user", "msg": ["Here is what came back:", photo]},
660+
("listing", listing),
661+
],
662+
questions={
663+
"same": Noul(instructions=["Is it the item in this photo?", listing]),
664+
"condition": Choice(criteria={"new": photo, "used": None}),
665+
"damage": Score(criteria=["none", {"example": [photo]}]),
666+
},
667+
)
668+
669+
sentry_sdk.flush()
670+
spans = [item.payload for item in items]
671+
(span,) = spans
672+
673+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
674+
{
675+
"type": "evaluation",
676+
"state": [
677+
{"role": "user", "msg": ["Here is what came back:", str(photo)]},
678+
["listing", str(listing)],
679+
],
680+
"questions": {
681+
"same": {
682+
"type": "noul",
683+
"instructions": ["Is it the item in this photo?", str(listing)],
684+
},
685+
"condition": {
686+
"type": "choice",
687+
"criteria": {"new": str(photo), "used": None},
688+
},
689+
"damage": {
690+
"type": "score",
691+
"criteria": ["none", {"example": [str(photo)]}],
692+
},
693+
},
694+
}
695+
]
696+
697+
698+
@pytest.mark.asyncio
699+
@pytest.mark.skipif(
700+
Image is None, reason="Image was added in version v0.7.5 of typesafe-sdk"
701+
)
702+
async def test_system_one_async_with_inline_image(
703+
sentry_init,
704+
capture_items,
705+
typesafe_response,
706+
):
707+
sentry_init(
708+
integrations=[TypeSafeIntegration()],
709+
traces_sample_rate=1.0,
710+
send_default_pii=True,
711+
trace_lifecycle="stream",
712+
)
713+
714+
photo = Image(b"\x89PNG\r\n\x1a\nimage data")
715+
listing = Image(b"\xff\xd8\xffimage data")
716+
717+
client = AsyncTypeSafeClient(api_key="z")
718+
items = capture_items("span")
719+
720+
with mock.patch.object(
721+
AsyncTypeSafeClient,
722+
"_request",
723+
new_callable=mock.AsyncMock,
724+
return_value=typesafe_response,
725+
), sentry_sdk.start_transaction(name="typesafe"):
726+
await client.system_one(
727+
state=[
728+
{"role": "user", "msg": ["Here is what came back:", photo]},
729+
("listing", listing),
730+
],
731+
questions={
732+
"same": Noul(instructions=["Is it the item in this photo?", listing]),
733+
"condition": Choice(criteria={"new": photo, "used": None}),
734+
"damage": Score(criteria=["none", {"example": [photo]}]),
735+
},
736+
)
737+
738+
sentry_sdk.flush()
739+
spans = [item.payload for item in items]
740+
(span,) = spans
741+
742+
assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
743+
{
744+
"type": "evaluation",
745+
"state": [
746+
{"role": "user", "msg": ["Here is what came back:", str(photo)]},
747+
["listing", str(listing)],
748+
],
749+
"questions": {
750+
"same": {
751+
"type": "noul",
752+
"instructions": ["Is it the item in this photo?", str(listing)],
753+
},
754+
"condition": {
755+
"type": "choice",
756+
"criteria": {"new": str(photo), "used": None},
757+
},
758+
"damage": {
759+
"type": "score",
760+
"criteria": ["none", {"example": [str(photo)]}],
761+
},
762+
},
763+
}
764+
]

0 commit comments

Comments
 (0)