feat: support multimodal input in agent chat

cognite/client/data_classes/agents/__init__.py

@@ -35,6 +35,7 @@
     ClientToolAction,
     ClientToolCall,
     ClientToolResult,
+    ImageContent,
     Message,
     MessageContent,
     MessageList,
@@ -73,6 +74,7 @@
     "ClientToolCall",
     "ClientToolResult",
     "DataModelInfo",
+    "ImageContent",
     "InstanceSpaces",
     "Message",
     "MessageContent",

cognite/client/data_classes/agents/chat.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
+import base64
 import json
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Any, ClassVar, Literal
 
 from cognite.client.data_classes._base import CogniteResource, CogniteResourceList
 from cognite.client.utils._text import convert_all_keys_to_camel_case
 
+_ALLOWED_IMAGE_MEDIA_TYPES = frozenset({"image/png", "image/jpeg", "image/webp"})
+_MEDIA_TYPE_BY_SUFFIX = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
+_MAX_IMAGE_DATA_LENGTH = 6_990_508
+
 
 @dataclass
 class MessageContent(CogniteResource, ABC):
@@ -34,6 +41,26 @@ def _load_content(cls, data: dict[str, Any]) -> MessageContent:
         """Create a concrete content instance from raw data."""
         ...
 
+    @classmethod
+    def _load_content_value(
+        cls, data: dict[str, Any] | list[dict[str, Any]] | None
+    ) -> MessageContent | list[MessageContent] | None:
+        """Load a single content object or a list of content parts."""
+        if data is None:
+            return None
+        if isinstance(data, list):
+            return [cls._load(item) for item in data]
+        return cls._load(data)
+
+    @staticmethod
+    def _dump_content_value(
+        content: MessageContent | list[MessageContent], camel_case: bool = True
+    ) -> dict[str, Any] | list[dict[str, Any]]:
+        """Dump a single content object or a list of content parts."""
+        if isinstance(content, list):
+            return [item.dump(camel_case=camel_case) for item in content]
+        return content.dump(camel_case=camel_case)
+
 
 @dataclass
 class TextContent(MessageContent):
@@ -51,6 +78,67 @@ def _load_content(cls, data: dict[str, Any]) -> TextContent:
         return cls(text=data["text"])
 
 
+@dataclass
+class ImageContent(MessageContent):
+    """Image content for messages.
+
+    Args:
+        data (str): Base64-encoded image data.
+        media_type (str): The MIME type of the image. Must be ``image/png``, ``image/jpeg``, or ``image/webp``.
+    """
+
+    _type: ClassVar[str] = "image"
+    data: str = ""
+    media_type: str = ""
+
+    @classmethod
+    def _load_content(cls, data: dict[str, Any]) -> ImageContent:
+        return cls(data=data["data"], media_type=data["mediaType"])
+
+    @classmethod
+    def from_bytes(cls, data: bytes, media_type: str) -> ImageContent:
+        """Create image content from raw image bytes.
+
+        Args:
+            data (bytes): Raw image bytes.
+            media_type (str): The MIME type of the image. Must be ``image/png``, ``image/jpeg``, or ``image/webp``.
+
+        Returns:
+            ImageContent: The image content ready to send to an agent.
+        """
+        if media_type not in _ALLOWED_IMAGE_MEDIA_TYPES:
+            raise ValueError(
+                f"Unsupported media type {media_type!r}. "
+                f"Allowed values: {', '.join(sorted(_ALLOWED_IMAGE_MEDIA_TYPES))}."
+            )
+        encoded = base64.b64encode(data).decode()
+        if len(encoded) > _MAX_IMAGE_DATA_LENGTH:
+            raise ValueError(
+                f"Encoded image data exceeds the maximum allowed length of {_MAX_IMAGE_DATA_LENGTH} characters."
+            )
+        return cls(data=encoded, media_type=media_type)
+
+    @classmethod
+    def from_file(cls, path: str | Path, media_type: str | None = None) -> ImageContent:
+        """Create image content from an image file on disk.
+
+        Args:
+            path (str | Path): Path to the image file.
+            media_type (str | None): The MIME type of the image. If not provided, it is inferred from the file suffix.
+
+        Returns:
+            ImageContent: The image content ready to send to an agent.
+        """
+        file_path = Path(path)
+        resolved_media_type = media_type or _MEDIA_TYPE_BY_SUFFIX.get(file_path.suffix.lower())
+        if resolved_media_type is None:
+            raise ValueError(
+                f"Could not infer media type from file suffix {file_path.suffix!r}. "
+                f"Provide media_type explicitly or use one of: {', '.join(sorted(_MEDIA_TYPE_BY_SUFFIX))}."
+            )
+        return cls.from_bytes(file_path.read_bytes(), resolved_media_type)
+
+
 @dataclass
 class UnknownContent(MessageContent):
     """Unknown content type for forward compatibility.
@@ -326,29 +414,36 @@ class Message(CogniteResource):
     """A message to send to an agent.
 
     Args:
-        content (str | MessageContent): The message content. If a string is provided, it will be converted to TextContent.
+        content (str | MessageContent | Sequence[MessageContent]): The message content. If a string is provided,
+            it will be converted to TextContent. Use a sequence of content parts to send multimodal messages with images.
         role (Literal['user']): The role of the message sender. Defaults to "user".
     """
 
-    content: MessageContent
+    content: MessageContent | list[MessageContent]
     role: Literal["user"] = "user"
 
-    def __init__(self, content: str | MessageContent, role: Literal["user"] = "user") -> None:
+    def __init__(
+        self, content: str | MessageContent | Sequence[MessageContent], role: Literal["user"] = "user"
+    ) -> None:
         if isinstance(content, str):
             self.content = TextContent(text=content)
-        else:
+        elif isinstance(content, MessageContent):
             self.content = content
+        else:
+            self.content = list(content)
         self.role = role
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         return {
-            "content": self.content.dump(camel_case=camel_case),
+            "content": MessageContent._dump_content_value(self.content, camel_case=camel_case),
             "role": self.role,
         }
 
     @classmethod
     def _load(cls, data: dict[str, Any]) -> Message:
-        content = MessageContent._load(data["content"])
+        content = MessageContent._load_content_value(data["content"])
+        if content is None:
+            raise ValueError("Message content is required.")
         return cls(content=content, role=data["role"])
 
 
@@ -597,14 +692,14 @@ class AgentMessage(CogniteResource):
     """A message from an agent.
 
     Args:
-        content (MessageContent | None): The message content.
+        content (MessageContent | list[MessageContent] | None): The message content.
         data (list[AgentDataItem] | None): Data items in the response.
         reasoning (list[AgentReasoningItem] | None): Reasoning items in the response.
         actions (list[ActionCall] | None): Action calls requested by the agent.
         role (Literal['agent']): The role of the message sender.
     """
 
-    content: MessageContent | None = None
+    content: MessageContent | list[MessageContent] | None = None
     data: list[AgentDataItem] | None = None
     reasoning: list[AgentReasoningItem] | None = None
     actions: list[ActionCall] | None = None
@@ -613,7 +708,7 @@ class AgentMessage(CogniteResource):
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         result: dict[str, Any] = {"role": self.role}
         if self.content is not None:
-            result["content"] = self.content.dump(camel_case=camel_case)
+            result["content"] = MessageContent._dump_content_value(self.content, camel_case=camel_case)
         if self.data is not None:
             result["data"] = [item.dump(camel_case=camel_case) for item in self.data]
         if self.reasoning is not None:
@@ -628,7 +723,7 @@ def _load(cls, data: dict[str, Any]) -> AgentMessage:
         reasoning_items = [AgentReasoningItem._load(item) for item in data.get("reasoning", [])]
         action_calls = [ActionCall._load(item) for item in data.get("actions", [])]
         return cls(
-            content=MessageContent._load_if(data.get("content")),
+            content=MessageContent._load_content_value(data.get("content")),
             data=data_items or None,
             reasoning=reasoning_items or None,
             actions=action_calls or None,
@@ -678,10 +773,17 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
     @property
     def text(self) -> str | None:
         """Get the text content from the first message with text content."""
-        if self.messages:
-            for message in self.messages:
-                if message.content and isinstance(message.content, TextContent):
-                    return message.content.text
+        if not self.messages:
+            return None
+        for message in self.messages:
+            if message.content is None:
+                continue
+            if isinstance(message.content, list):
+                for part in message.content:
+                    if isinstance(part, TextContent):
+                        return part.text
+            elif isinstance(message.content, TextContent):
+                return message.content.text
         return None
 
     @property