kosong.message

11class MergeableMixin:
12    def merge_in_place(self, other: Any) -> bool:
13        """Merge the other part into the current part. Return True if the merge is successful."""
14        return False

17class ContentPart(BaseModel, ABC, MergeableMixin):
18    """
19    A part of a message content.
20
21    This is the abstract base class for all supported content parts. Subclasses must define a `type`
22    field of type `str` and optional other fields specific to the content part.
23
24    For Kosong users, you typically do not need to subclass this directly. Instead, use the provided
25    subclasses like `TextPart`, `ThinkPart`, `ImageURLPart`, etc. Unless you are implementing custom
26    `ChatProvider`s that supports new content part types.
27    """
28
29    __content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {}
30
31    type: str
32    ...  # to be added by subclasses
33
34    def __init_subclass__(cls, **kwargs: Any) -> None:
35        super().__init_subclass__(**kwargs)
36
37        invalid_subclass_error_msg = (
38            f"ContentPart subclass {cls.__name__} must have a `type` field of type `str`"
39        )
40
41        type_value = getattr(cls, "type", None)
42        if type_value is None or not isinstance(type_value, str):
43            raise ValueError(invalid_subclass_error_msg)
44
45        cls.__content_part_registry[type_value] = cls
46
47    @classmethod
48    def __get_pydantic_core_schema__(
49        cls, source_type: Any, handler: GetCoreSchemaHandler
50    ) -> core_schema.CoreSchema:
51        # If we're dealing with the base ContentPart class, use custom validation
52        if cls.__name__ == "ContentPart":
53
54            def validate_content_part(value: Any) -> Any:
55                # if it's already an instance of a ContentPart subclass, return it
56                if hasattr(value, "__class__") and issubclass(value.__class__, cls):
57                    return value
58
59                # if it's a dict with a type field, dispatch to the appropriate subclass
60                if isinstance(value, dict) and "type" in value:
61                    type_value: Any | None = cast(dict[str, Any], value).get("type")
62                    if not isinstance(type_value, str):
63                        raise ValueError(f"Cannot validate {value} as ContentPart")
64                    target_class = cls.__content_part_registry[type_value]
65                    return target_class.model_validate(value)
66
67                raise ValueError(f"Cannot validate {value} as ContentPart")
68
69            return core_schema.no_info_plain_validator_function(validate_content_part)
70
71        # for subclasses, use the default schema
72        return handler(source_type)

75class TextPart(ContentPart):
76    """
77    >>> TextPart(text="Hello, world!").model_dump()
78    {'type': 'text', 'text': 'Hello, world!'}
79    """
80
81    type: str = "text"
82    text: str
83
84    @override
85    def merge_in_place(self, other: Any) -> bool:
86        if not isinstance(other, TextPart):
87            return False
88        self.text += other.text
89        return True

 92class ThinkPart(ContentPart):
 93    """
 94    >>> ThinkPart(think="I think I need to think about this.").model_dump()
 95    {'type': 'think', 'think': 'I think I need to think about this.', 'encrypted': None}
 96    """
 97
 98    type: str = "think"
 99    think: str
100    encrypted: str | None = None
101    """Encrypted thinking content, or signature."""
102
103    @override
104    def merge_in_place(self, other: Any) -> bool:
105        if not isinstance(other, ThinkPart):
106            return False
107        if self.encrypted:
108            return False
109        self.think += other.think
110        if other.encrypted:
111            self.encrypted = other.encrypted
112        return True

115class ImageURLPart(ContentPart):
116    """
117    >>> ImageURLPart(
118    ...     image_url=ImageURLPart.ImageURL(url="https://example.com/image.png")
119    ... ).model_dump()
120    {'type': 'image_url', 'image_url': {'url': 'https://example.com/image.png', 'id': None}}
121    """
122
123    class ImageURL(BaseModel):
124        """Image URL payload."""
125
126        url: str
127        """The URL of the image, can be data URI scheme like `data:image/png;base64,...`."""
128        id: str | None = None
129        """The ID of the image, to allow LLMs to distinguish different images."""
130
131    type: str = "image_url"
132    image_url: ImageURL

123    class ImageURL(BaseModel):
124        """Image URL payload."""
125
126        url: str
127        """The URL of the image, can be data URI scheme like `data:image/png;base64,...`."""
128        id: str | None = None
129        """The ID of the image, to allow LLMs to distinguish different images."""

135class AudioURLPart(ContentPart):
136    """
137    >>> AudioURLPart(
138    ...     audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")
139    ... ).model_dump()
140    {'type': 'audio_url', 'audio_url': {'url': 'https://example.com/audio.mp3', 'id': None}}
141    """
142
143    class AudioURL(BaseModel):
144        """Audio URL payload."""
145
146        url: str
147        """The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`."""
148        id: str | None = None
149        """The ID of the audio, to allow LLMs to distinguish different audios."""
150
151    type: str = "audio_url"
152    audio_url: AudioURL

143    class AudioURL(BaseModel):
144        """Audio URL payload."""
145
146        url: str
147        """The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`."""
148        id: str | None = None
149        """The ID of the audio, to allow LLMs to distinguish different audios."""

155class VideoURLPart(ContentPart):
156    """
157    >>> VideoURLPart(
158    ...     video_url=VideoURLPart.VideoURL(url="https://example.com/video.mp4")
159    ... ).model_dump()
160    {'type': 'video_url', 'video_url': {'url': 'https://example.com/video.mp4', 'id': None}}
161    """
162
163    class VideoURL(BaseModel):
164        """Video URL payload."""
165
166        url: str
167        """The URL of the video, can be data URI scheme like `data:video/mp4;base64,...`."""
168        id: str | None = None
169        """The ID of the video, to allow LLMs to distinguish different videos."""
170
171    type: str = "video_url"
172    video_url: VideoURL

163    class VideoURL(BaseModel):
164        """Video URL payload."""
165
166        url: str
167        """The URL of the video, can be data URI scheme like `data:video/mp4;base64,...`."""
168        id: str | None = None
169        """The ID of the video, to allow LLMs to distinguish different videos."""

175class ToolCall(BaseModel, MergeableMixin):
176    """
177    A tool call requested by the assistant.
178
179    >>> ToolCall(
180    ...     id="123",
181    ...     function=ToolCall.FunctionBody(name="function", arguments="{}"),
182    ... ).model_dump(exclude_none=True)
183    {'type': 'function', 'id': '123', 'function': {'name': 'function', 'arguments': '{}'}}
184    """
185
186    class FunctionBody(BaseModel):
187        """Tool call function body."""
188
189        name: str
190        """The name of the tool to be called."""
191        arguments: str | None
192        """Arguments of the tool call in JSON string format."""
193
194    type: Literal["function"] = "function"
195
196    id: str
197    """The ID of the tool call."""
198    function: FunctionBody
199    """The function body of the tool call."""
200    extras: dict[str, JsonType] | None = None
201    """Extra information about the tool call."""
202
203    @override
204    def merge_in_place(self, other: Any) -> bool:
205        if not isinstance(other, ToolCallPart):
206            return False
207        if self.function.arguments is None:
208            self.function.arguments = other.arguments_part
209        else:
210            self.function.arguments += other.arguments_part or ""
211        return True

186    class FunctionBody(BaseModel):
187        """Tool call function body."""
188
189        name: str
190        """The name of the tool to be called."""
191        arguments: str | None
192        """Arguments of the tool call in JSON string format."""

214class ToolCallPart(BaseModel, MergeableMixin):
215    """A part of the tool call."""
216
217    arguments_part: str | None = None
218    """A part of the arguments of the tool call."""
219
220    @override
221    def merge_in_place(self, other: Any) -> bool:
222        if not isinstance(other, ToolCallPart):
223            return False
224        if self.arguments_part is None:
225            self.arguments_part = other.arguments_part
226        else:
227            self.arguments_part += other.arguments_part or ""
228        return True

244class Message(BaseModel):
245    """A message in a conversation."""
246
247    role: Role
248    """The role of the message sender."""
249
250    name: str | None = None
251
252    content: list[ContentPart]
253    """
254    The content of the message.
255    Empty list `[]` will be interpreted as no content.
256    """
257
258    tool_calls: list[ToolCall] | None = None
259    """Tool calls requested by the assistant in this message."""
260
261    tool_call_id: str | None = None
262    """The ID of the tool call if this message is a tool response."""
263
264    partial: bool | None = None
265
266    @field_serializer("content")
267    def _serialize_content(self, content: list[ContentPart]) -> str | list[dict[str, Any]] | None:
268        if len(content) == 1 and isinstance(content[0], TextPart):
269            return content[0].text
270        return [part.model_dump() for part in content]
271
272    @field_validator("content", mode="before")
273    @classmethod
274    def _coerce_none_content(cls, value: Any) -> Any:
275        if value is None:
276            return []
277        if isinstance(value, str):
278            return [TextPart(text=value)]
279        return value
280
281    def __init__(
282        self,
283        *,
284        role: Role,
285        content: list[ContentPart] | ContentPart | str,
286        tool_calls: list[ToolCall] | None = None,
287        tool_call_id: str | None = None,
288        **data: Any,
289    ) -> None:
290        if isinstance(content, str):
291            content = [TextPart(text=content)]
292        elif isinstance(content, ContentPart):
293            content = [content]
294        super().__init__(
295            role=role,
296            content=content,
297            tool_calls=tool_calls,
298            tool_call_id=tool_call_id,
299            **data,
300        )
301
302    def extract_text(self, sep: str = "") -> str:
303        """Extract and concatenate all text parts in the message content."""
304        return sep.join(part.text for part in self.content if isinstance(part, TextPart))