diff --git a/cognite/client/data_classes/agents/__init__.py b/cognite/client/data_classes/agents/__init__.py index b03eecc9cb..06aba408a0 100644 --- a/cognite/client/data_classes/agents/__init__.py +++ b/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", diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py index a3a1e34be1..1a3705f524 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -1,13 +1,19 @@ 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"} + @dataclass class MessageContent(CogniteResource, ABC): @@ -34,6 +40,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 +77,63 @@ 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() + 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 +409,45 @@ 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 | str]): 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. Strings in a sequence are also converted to TextContent. 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 | str], role: Literal["user"] = "user" + ) -> None: if isinstance(content, str): self.content = TextContent(text=content) - else: + elif isinstance(content, MessageContent): self.content = content + else: + parts: list[MessageContent] = [] + for item in content: + if isinstance(item, str): + parts.append(TextContent(text=item)) + elif isinstance(item, MessageContent): + parts.append(item) + else: + raise TypeError(f"Expected str or MessageContent, got {type(item).__name__}") + self.content = parts 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 +696,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 +712,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 +727,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 +777,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 diff --git a/tests/tests_unit/test_api/test_agents_chat.py b/tests/tests_unit/test_api/test_agents_chat.py index 445bec3730..0967df44c8 100644 --- a/tests/tests_unit/test_api/test_agents_chat.py +++ b/tests/tests_unit/test_api/test_agents_chat.py @@ -1,5 +1,7 @@ from __future__ import annotations +import base64 +from pathlib import Path from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock @@ -7,7 +9,7 @@ from pytest_httpx import HTTPXMock from cognite.client import AsyncCogniteClient, CogniteClient -from cognite.client.data_classes.agents import Message +from cognite.client.data_classes.agents import ImageContent, Message from cognite.client.data_classes.agents.chat import ( AgentChatResponse, AgentDataItem, @@ -222,3 +224,178 @@ def test_message_with_explicit_content(self) -> None: assert msg.content is content assert isinstance(msg.content, TextContent) assert msg.content.text == "Hello world" + + +class TestImageContent: + def test_from_bytes(self) -> None: + image_bytes = b"fake-image-bytes" + content = ImageContent.from_bytes(image_bytes, media_type="image/png") + + assert content.data == base64.b64encode(image_bytes).decode() + assert content.media_type == "image/png" + assert content.dump() == { + "data": base64.b64encode(image_bytes).decode(), + "mediaType": "image/png", + "type": "image", + } + + def test_from_bytes_invalid_media_type(self) -> None: + with pytest.raises(ValueError, match="Unsupported media type"): + ImageContent.from_bytes(b"data", media_type="image/gif") + + def test_from_file(self, tmp_path: Path) -> None: + image_bytes = b"fake-image-bytes" + image_path = tmp_path / "diagram.png" + image_path.write_bytes(image_bytes) + + content = ImageContent.from_file(image_path) + + assert content.data == base64.b64encode(image_bytes).decode() + assert content.media_type == "image/png" + + def test_from_file_with_explicit_media_type(self, tmp_path: Path) -> None: + image_bytes = b"fake-image-bytes" + image_path = tmp_path / "diagram.bin" + image_path.write_bytes(image_bytes) + + content = ImageContent.from_file(image_path, media_type="image/webp") + + assert content.media_type == "image/webp" + + def test_from_file_unknown_suffix(self, tmp_path: Path) -> None: + image_path = tmp_path / "diagram.bin" + image_path.write_bytes(b"data") + + with pytest.raises(ValueError, match="Could not infer media type"): + ImageContent.from_file(image_path) + + def test_load_from_api_response(self) -> None: + data = {"type": "image", "data": "aGVsbG8=", "mediaType": "image/jpeg"} + content = ImageContent._load(data) + + assert isinstance(content, ImageContent) + assert content.data == "aGVsbG8=" + assert content.media_type == "image/jpeg" + + +class TestMultimodalMessage: + def test_message_with_content_parts_dump(self) -> None: + image_bytes = b"fake-image-bytes" + message = Message( + [ + TextContent(text="What's in this image?"), + ImageContent.from_bytes(image_bytes, media_type="image/png"), + ] + ) + + dumped = message.dump() + + assert dumped == { + "role": "user", + "content": [ + {"text": "What's in this image?", "type": "text"}, + { + "data": base64.b64encode(image_bytes).decode(), + "mediaType": "image/png", + "type": "image", + }, + ], + } + + def test_text_only_message_still_dumps_single_object(self) -> None: + message = Message("Hello world") + + assert message.dump() == { + "role": "user", + "content": {"text": "Hello world", "type": "text"}, + } + + def test_message_with_string_content_parts(self) -> None: + image_bytes = b"fake-image-bytes" + message = Message( + [ + "What's in this image?", + ImageContent.from_bytes(image_bytes, media_type="image/png"), + ] + ) + + assert isinstance(message.content, list) + assert isinstance(message.content[0], TextContent) + assert message.content[0].text == "What's in this image?" + assert isinstance(message.content[1], ImageContent) + + def test_message_with_invalid_content_part_raises_type_error(self) -> None: + with pytest.raises(TypeError, match="Expected str or MessageContent, got int"): + Message([123]) # type: ignore[list-item] + + def test_chat_with_multimodal_message( + self, + httpx_mock: HTTPXMock, + cognite_client: CogniteClient, + async_client: AsyncCogniteClient, + chat_response_body: dict, + ) -> None: + httpx_mock.add_response( + method="POST", + url=get_url(async_client.agents, async_client.agents._RESOURCE_PATH + "/chat"), + status_code=200, + json=chat_response_body, + ) + + image_bytes = b"fake-image-bytes" + cognite_client.agents.chat( + agent_external_id="my_agent", + messages=Message( + [ + TextContent(text="Describe this image"), + ImageContent.from_bytes(image_bytes, media_type="image/jpeg"), + ] + ), + ) + + request = httpx_mock.get_requests()[0] + payload = jsgz_load(request.content) + assert payload == { + "agentExternalId": "my_agent", + "messages": [ + { + "role": "user", + "content": [ + {"text": "Describe this image", "type": "text"}, + { + "data": base64.b64encode(image_bytes).decode(), + "mediaType": "image/jpeg", + "type": "image", + }, + ], + } + ], + } + + def test_load_agent_message_with_content_parts(self) -> None: + response = AgentChatResponse._load( + { + "agentExternalId": "my_agent", + "response": { + "type": "result", + "messages": [ + { + "role": "agent", + "content": [ + {"text": "I see a pump in the image.", "type": "text"}, + {"data": "aGVsbG8=", "mediaType": "image/png", "type": "image"}, + ], + } + ], + }, + } + ) + + assert response.text == "I see a pump in the image." + content = response.messages[0].content + assert isinstance(content, list) + assert isinstance(content[0], TextContent) + assert isinstance(content[1], ImageContent) + assert content[0].text == "I see a pump in the image." + assert content[1].data == "aGVsbG8=" + assert content[1].media_type == "image/png"