Use proto descriptor for typename (#5346)

Closes #5302
This commit is contained in:
Jack Gerrits 2025-02-04 14:18:16 -05:00 committed by GitHub
parent 6454e3fece
commit 40d74a32a1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 1 deletions

View File

@ -192,6 +192,13 @@ class UnknownPayload:
def _type_name(cls: type[Any] | Any) -> str:
# If cls is a protobuf, then we need to determine the descriptor
if isinstance(cls, type):
if issubclass(cls, Message):
return cast(str, cls.DESCRIPTOR.full_name)
elif isinstance(cls, Message):
return cast(str, cls.DESCRIPTOR.full_name)
if isinstance(cls, type):
return cls.__name__
else:

View File

@ -93,7 +93,7 @@ def test_proto() -> None:
message = ProtoMessage(message="hello")
name = serde.type_name(message)
data = serde.serialize(message, type_name=name, data_content_type=PROTOBUF_DATA_CONTENT_TYPE)
assert name == "ProtoMessage"
assert name == "agents.ProtoMessage"
deserialized = serde.deserialize(data, type_name=name, data_content_type=PROTOBUF_DATA_CONTENT_TYPE)
assert deserialized.message == message.message
@ -186,3 +186,17 @@ def test_image_type() -> None:
assert deserialized.image.image.size == (100, 100)
assert deserialized.image.image.mode == "RGB"
assert deserialized.image.image == image.image
def test_type_name_for_protos() -> None:
type_name = SerializationRegistry().type_name(ProtoMessage())
assert type_name == "agents.ProtoMessage"
type_name = SerializationRegistry().type_name(ProtoMessage)
assert type_name == "agents.ProtoMessage"
type_name = SerializationRegistry().type_name(NestingProtoMessage())
assert type_name == "agents.NestingProtoMessage"
type_name = SerializationRegistry().type_name(NestingProtoMessage)
assert type_name == "agents.NestingProtoMessage"