Move PIL to optional import block

autogen/messages/agent_messages.py

@@ -7,14 +7,17 @@
 from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
 from uuid import UUID
 
-from PIL.Image import Image
 from pydantic import BaseModel, field_validator
 from termcolor import colored
 
 from ..code_utils import content_str
+from ..import_utils import optional_import_block, require_optional_import
 from ..oai.client import OpenAIWrapper
 from .base_message import BaseMessage, wrap_message
 
+with optional_import_block():
+    from PIL.Image import Image
+
 if TYPE_CHECKING:
     from ..agentchat.agent import Agent
     from ..coding.base import CodeBlock
@@ -188,6 +191,7 @@ def print(self, f: Optional[Callable[..., Any]] = None) -> None:
         f("\n", "-" * 80, flush=True, sep="")
 
 
+@require_optional_import("PIL", "unknown", except_for="print")
 @wrap_message
 class TextMessage(BasePrintReceivedMessage):
     content: Optional[Union[str, int, float, bool, list[dict[str, Union[str, dict[str, Any]]]]]] = None  # type: ignore [assignment]

test/messages/test_agent_messages.py

@@ -6,12 +6,12 @@
 from unittest.mock import MagicMock, _Call, call
 from uuid import UUID
 
-import PIL
 import pytest
 import termcolor.termcolor
 
 from autogen.agentchat.conversable_agent import ConversableAgent
 from autogen.coding.base import CodeBlock
+from autogen.import_utils import optional_import_block, skip_on_missing_imports
 from autogen.messages.agent_messages import (
     ClearAgentsHistoryMessage,
     ClearConversableAgentHistoryMessage,
@@ -42,6 +42,9 @@
     create_received_message_model,
 )
 
+with optional_import_block():
+    import PIL
+
 
 @pytest.fixture(autouse=True)
 def enable_color_in_tests(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -415,6 +418,7 @@ def test_print_context_lambda_message(
 
         assert mock.call_args_list == expected_call_args_list
 
+    @skip_on_missing_imports("PIL", "unknown")
     def test_serialization(self) -> None:
         image = PIL.Image.new(mode="RGB", size=(200, 200))
         content = [