Fix typing issues

Author: cpsievert

shiny/ui/_chat.py

@@ -540,7 +540,7 @@ async def _append_message(
                 self._current_stream_message = ""
 
         msg = await self._transform_message(
-            msg, chunk=chunk, chunk_content=chunk_content
+            msg, chunk=chunk, chunk_content=str(chunk_content)
         )
         if msg is None:
             return
@@ -776,11 +776,11 @@ async def _transform_message(
         key = res["transform_key"]
 
         if message["role"] == "user" and self._transform_user is not None:
-            content = await self._transform_user(message["content"])
+            content = await self._transform_user(str(message["content"]))
 
         elif message["role"] == "assistant" and self._transform_assistant is not None:
             content = await self._transform_assistant(
-                message["content"],
+                str(message["content"]),
                 chunk_content or "",
                 chunk == "end" or chunk is False,
             )
@@ -790,7 +790,7 @@ async def _transform_message(
         if content is None:
             return None
 
-        res[key] = content
+        res[key] = content  # type: ignore
 
         return res
 
@@ -950,7 +950,7 @@ def user_input(self, transform: bool = False) -> str | None:
         if msg is None:
             return None
         key = "content_server" if transform else "content_client"
-        return msg[key]
+        return str(msg[key])
 
     def _user_input(self) -> str:
         id = self.user_input_id
@@ -1097,7 +1097,7 @@ def as_transformed_message(message: ChatMessage) -> TransformedMessage:
 
     return TransformedMessage(
         content_client=message["content"],
-        content_server=message["content"],
+        content_server=str(message["content"]),
         role=message["role"],
         transform_key=transform_key,
         pre_transform_key=pre_transform_key,

shiny/ui/_chat_types.py

@@ -2,19 +2,21 @@
 
 from typing import Literal, TypedDict
 
+from htmltools import HTML
+
 Role = Literal["assistant", "user", "system"]
 
 
 # TODO: content should probably be [{"type": "text", "content": "..."}, {"type": "image", ...}]
 # in order to support multiple content types...
 class ChatMessage(TypedDict):
-    content: str
+    content: str | HTML
     role: Role
 
 
 # A message once transformed have been applied
 class TransformedMessage(TypedDict):
-    content_client: str
+    content_client: str | HTML
     content_server: str
     role: Role
     transform_key: Literal["content_client", "content_server"]

tests/playwright/shiny/components/chat/transform_assistant/app.py

@@ -12,7 +12,7 @@
 
 # TODO: test with append_message_stream() as well
 @chat.transform_assistant_response
-def transform(content: str) -> str:
+def transform(content: str) -> str | ui.HTML:
     if content == "return HTML":
         return ui.HTML(f"<b>Transformed response</b>: {content}")
     else: