Use a classmethod

Author: cpsievert

shiny/ui/_chat.py

@@ -251,7 +251,10 @@ async def _on_user_input():
                 else:
                     # A transformed value of None is a special signal to suspend input
                     # handling (i.e., don't generate a response)
-                    self._store_message(msg.as_transformed_message(), index=n_pre)
+                    self._store_message(
+                        TransformedMessage.from_chat_message(msg),
+                        index=n_pre,
+                    )
                     await self._remove_loading_message()
                     self._suspend_input_handler = True
 
@@ -935,7 +938,7 @@ async def _transform_message(
         chunk: ChunkOption = False,
         chunk_content: str | None = None,
     ) -> TransformedMessage | None:
-        res = message.as_transformed_message()
+        res = TransformedMessage.from_chat_message(message)
         key = res.transform_key
 
         if message.role == "user" and self._transform_user is not None:

shiny/ui/_chat_types.py

@@ -38,23 +38,6 @@ def __init__(
         self.content = content
         self.html_deps = deps
 
-    def as_transformed_message(self) -> "TransformedMessage":
-        if self.role == "user":
-            transform_key = "content_server"
-            pre_transform_key = "content_client"
-        else:
-            transform_key = "content_client"
-            pre_transform_key = "content_server"
-
-        return TransformedMessage(
-            content_client=self.content,
-            content_server=self.content,
-            role=self.role,
-            transform_key=transform_key,
-            pre_transform_key=pre_transform_key,
-            html_deps=self.html_deps,
-        )
-
 
 # A message once transformed have been applied
 @dataclass
@@ -66,6 +49,24 @@ class TransformedMessage:
     pre_transform_key: Literal["content_client", "content_server"]
     html_deps: list[dict[str, str]] | None = None
 
+    @classmethod
+    def from_chat_message(cls, message: ChatUIMessage) -> "TransformedMessage":
+        if message.role == "user":
+            transform_key = "content_server"
+            pre_transform_key = "content_client"
+        else:
+            transform_key = "content_client"
+            pre_transform_key = "content_server"
+
+        return TransformedMessage(
+            content_client=message.content,
+            content_server=message.content,
+            role=message.role,
+            transform_key=transform_key,
+            pre_transform_key=pre_transform_key,
+            html_deps=message.html_deps,
+        )
+
 
 # A message that can be sent to the client
 class ClientMessage(TypedDict):

tests/pytest/test_chat.py

@@ -13,7 +13,7 @@
 from shiny.types import MISSING
 from shiny.ui import Chat
 from shiny.ui._chat_normalize import normalize_message, normalize_message_chunk
-from shiny.ui._chat_types import ChatMessage, ChatUIMessage
+from shiny.ui._chat_types import ChatMessage, ChatUIMessage, Role, TransformedMessage
 
 # ----------------------------------------------------------------------
 # Helpers
@@ -42,6 +42,12 @@ def is_type_in_union(type: object, union: object) -> bool:
     return False
 
 
+def transformed_message(content: str, role: Role) -> TransformedMessage:
+    return TransformedMessage.from_chat_message(
+        ChatUIMessage(content=content, role=role)
+    )
+
+
 def test_chat_message_trimming():
     with session_context(test_session):
         chat = Chat(id="chat")
@@ -52,22 +58,19 @@ def generate_content(token_count: int) -> str:
             return " ".join(["foo" for _ in range(1, n)])
 
         msgs = (
-            ChatUIMessage(
-                content=generate_content(102), role="system"
-            ).as_transformed_message(),
+            transformed_message(
+                content=generate_content(102),
+                role="system",
+            ),
         )
 
         # Throws since system message is too long
         with pytest.raises(ValueError):
             chat._trim_messages(msgs, token_limits=(100, 0), format=MISSING)
 
         msgs = (
-            ChatUIMessage(
-                content=generate_content(100), role="system"
-            ).as_transformed_message(),
-            ChatUIMessage(
-                content=generate_content(2), role="user"
-            ).as_transformed_message(),
+            transformed_message(content=generate_content(100), role="system"),
+            transformed_message(content=generate_content(2), role="user"),
         )
 
         # Throws since only the system message fits
@@ -83,18 +86,18 @@ def generate_content(token_count: int) -> str:
         content3 = generate_content(2)
 
         msgs = (
-            ChatUIMessage(
+            transformed_message(
                 content=content1,
                 role="system",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content2,
                 role="user",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content3,
                 role="user",
-            ).as_transformed_message(),
+            ),
         )
 
         # Should discard the 1st user message
@@ -109,22 +112,22 @@ def generate_content(token_count: int) -> str:
         content4 = generate_content(2)
 
         msgs = (
-            ChatUIMessage(
+            transformed_message(
                 content=content1,
                 role="system",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content2,
                 role="user",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content3,
                 role="system",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content4,
                 role="user",
-            ).as_transformed_message(),
+            ),
         )
 
         # Should discard the 1st user message
@@ -137,14 +140,14 @@ def generate_content(token_count: int) -> str:
         content2 = generate_content(10)
 
         msgs = (
-            ChatUIMessage(
+            transformed_message(
                 content=content1,
                 role="assistant",
-            ).as_transformed_message(),
-            ChatUIMessage(
+            ),
+            transformed_message(
                 content=content2,
                 role="user",
-            ).as_transformed_message(),
+            ),
         )
 
         # Anthropic requires 1st message to be a user message