diff --git a/packages/kosong/src/kosong/message.py b/packages/kosong/src/kosong/message.py index a658f36992..5351c8d25f 100644 --- a/packages/kosong/src/kosong/message.py +++ b/packages/kosong/src/kosong/message.py @@ -1,7 +1,13 @@ from abc import ABC from typing import Any, ClassVar, Literal, cast, override -from pydantic import BaseModel, GetCoreSchemaHandler, field_serializer, field_validator +from pydantic import ( + BaseModel, + GetCoreSchemaHandler, + SerializerFunctionWrapHandler, + field_serializer, + field_validator, +) from pydantic_core import core_schema from kosong.utils.typing import JsonType @@ -262,11 +268,15 @@ class Message(BaseModel): partial: bool | None = None - @field_serializer("content") - def _serialize_content(self, content: list[ContentPart]) -> str | list[dict[str, Any]] | None: + @field_serializer("content", mode="wrap") + def _serialize_content( + self, + content: list[ContentPart], + serializer: SerializerFunctionWrapHandler, + ) -> str | list[dict[str, Any]] | None: if len(content) == 1 and isinstance(content[0], TextPart): return content[0].text - return [part.model_dump() for part in content] + return serializer(content) @field_validator("content", mode="before") @classmethod diff --git a/packages/kosong/tests/api_snapshot_tests/test_kimi.py b/packages/kosong/tests/api_snapshot_tests/test_kimi.py index 369f3c32f9..6e28409bdd 100644 --- a/packages/kosong/tests/api_snapshot_tests/test_kimi.py +++ b/packages/kosong/tests/api_snapshot_tests/test_kimi.py @@ -133,7 +133,6 @@ async def test_kimi_message_conversion(): "type": "image_url", "image_url": { "url": "https://example.com/image.png", - "id": None, }, }, ], @@ -204,7 +203,6 @@ async def test_kimi_message_conversion(): "type": "image_url", "image_url": { "url": "https://example.com/image.png", - "id": None, }, }, ], diff --git a/packages/kosong/tests/api_snapshot_tests/test_openai_legacy.py b/packages/kosong/tests/api_snapshot_tests/test_openai_legacy.py index 78114ff9ae..e773555a72 100644 --- a/packages/kosong/tests/api_snapshot_tests/test_openai_legacy.py +++ b/packages/kosong/tests/api_snapshot_tests/test_openai_legacy.py @@ -57,7 +57,6 @@ async def test_openai_legacy_message_conversion(): "type": "image_url", "image_url": { "url": "https://example.com/image.png", - "id": None, }, }, ], @@ -128,7 +127,6 @@ async def test_openai_legacy_message_conversion(): "type": "image_url", "image_url": { "url": "https://example.com/image.png", - "id": None, }, }, ], diff --git a/packages/kosong/tests/test_message.py b/packages/kosong/tests/test_message.py index 7dfa1056db..ba6a01f907 100644 --- a/packages/kosong/tests/test_message.py +++ b/packages/kosong/tests/test_message.py @@ -30,7 +30,7 @@ def test_message_with_single_part(): "content": [ { "type": "image_url", - "image_url": {"url": "https://example.com/image.png", "id": None}, + "image_url": {"url": "https://example.com/image.png"}, } ], } @@ -38,6 +38,46 @@ def test_message_with_single_part(): assert Message.model_validate(dumped) == message +def test_exclude_none_applies_to_nested_media_content_parts(): + message = Message( + role="user", + content=[ + ImageURLPart(image_url=ImageURLPart.ImageURL(url="https://example.com/image.png")), + AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")), + VideoURLPart( + video_url=VideoURLPart.VideoURL( + url="https://example.com/video.mp4", + id="video-1", + ) + ), + ], + ) + + dumped_with_none = message.model_dump() + assert dumped_with_none["content"][0]["image_url"]["id"] is None + + assert message.model_dump(exclude_none=True) == { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + }, + { + "type": "audio_url", + "audio_url": {"url": "https://example.com/audio.mp3"}, + }, + { + "type": "video_url", + "video_url": { + "url": "https://example.com/video.mp4", + "id": "video-1", + }, + }, + ], + } + + def test_message_with_tool_calls(): message = Message( role="assistant", @@ -107,22 +147,18 @@ def test_message_with_complex_content(): "role": "user", "content": [ {"type": "text", "text": "Hello, world!"}, - { - "type": "think", - "think": "I think I need to think about this.", - "encrypted": None, - }, + {"type": "think", "think": "I think I need to think about this."}, { "type": "image_url", - "image_url": {"url": "https://example.com/image.png", "id": None}, + "image_url": {"url": "https://example.com/image.png"}, }, { "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3", "id": None}, + "audio_url": {"url": "https://example.com/audio.mp3"}, }, { "type": "video_url", - "video_url": {"url": "https://example.com/video.mp4", "id": None}, + "video_url": {"url": "https://example.com/video.mp4"}, }, ], "tool_calls": [