Rename ChatMessage -> ChatMessageDict; ChatUIMessage -> ChatMessage

Author: cpsievert

shiny/ui/_chat.py

@@ -38,15 +38,15 @@
     as_provider_message,
 )
 from ._chat_tokenizer import TokenEncoding, TokenizersEncoding, get_default_tokenizer
-from ._chat_types import ChatMessage, ChatUIMessage, ClientMessage, TransformedMessage
+from ._chat_types import ChatMessage, ChatMessageDict, ClientMessage, TransformedMessage
 from ._html_deps_py_shiny import chat_deps
 from .fill import as_fill_item, as_fillable_container
 
 __all__ = (
     "Chat",
     "ChatExpress",
     "chat_ui",
-    "ChatMessage",
+    "ChatMessageDict",
 )
 
 
@@ -240,7 +240,7 @@ async def _init_chat():
             @reactive.effect(priority=9999)
             @reactive.event(self._user_input)
             async def _on_user_input():
-                msg = ChatUIMessage(content=self._user_input(), role="user")
+                msg = ChatMessage(content=self._user_input(), role="user")
                 # It's possible that during the transform, a message is appended, so get
                 # the length now, so we can insert the new message at the right index
                 n_pre = len(self._messages())
@@ -415,7 +415,7 @@ def messages(
         token_limits: tuple[int, int] | None = None,
         transform_user: Literal["all", "last", "none"] = "all",
         transform_assistant: bool = False,
-    ) -> tuple[ChatMessage, ...]: ...
+    ) -> tuple[ChatMessageDict, ...]: ...
 
     def messages(
         self,
@@ -424,7 +424,7 @@ def messages(
         token_limits: tuple[int, int] | None = None,
         transform_user: Literal["all", "last", "none"] = "all",
         transform_assistant: bool = False,
-    ) -> tuple[ChatMessage | ProviderMessage, ...]:
+    ) -> tuple[ChatMessageDict | ProviderMessage, ...]:
         """
         Reactively read chat messages
 
@@ -492,7 +492,7 @@ def messages(
         if token_limits is not None:
             messages = self._trim_messages(messages, token_limits, format)
 
-        res: list[ChatMessage | ProviderMessage] = []
+        res: list[ChatMessageDict | ProviderMessage] = []
         for i, m in enumerate(messages):
             transform = False
             if m.role == "assistant":
@@ -505,7 +505,7 @@ def messages(
                 m, "transform_key" if transform else "pre_transform_key"
             )
             content = getattr(m, content_key)
-            chat_msg = ChatMessage(content=str(content), role=m.role)
+            chat_msg = ChatMessageDict(content=str(content), role=m.role)
             if not isinstance(format, MISSING_TYPE):
                 chat_msg = as_provider_message(chat_msg, format)
             res.append(chat_msg)
@@ -745,7 +745,7 @@ async def _append_message_stream(
     ):
         id = _utils.private_random_id()
 
-        empty = ChatMessage(content="", role="assistant")
+        empty = ChatMessageDict(content="", role="assistant")
         await self._append_message(empty, chunk="start", stream_id=id, icon=icon)
 
         try:
@@ -934,7 +934,7 @@ async def _transform_wrapper(content: str, chunk: str, done: bool):
 
     async def _transform_message(
         self,
-        message: ChatUIMessage,
+        message: ChatMessage,
         chunk: ChunkOption = False,
         chunk_content: str | None = None,
     ) -> TransformedMessage | None:
@@ -1201,7 +1201,7 @@ class ChatExpress(Chat):
     def ui(
         self,
         *,
-        messages: Optional[Sequence[str | ChatMessage]] = None,
+        messages: Optional[Sequence[str | ChatMessageDict]] = None,
         placeholder: str = "Enter a message...",
         width: CssUnit = "min(680px, 100%)",
         height: CssUnit = "auto",
@@ -1251,7 +1251,7 @@ def ui(
 def chat_ui(
     id: str,
     *,
-    messages: Optional[Sequence[TagChild | ChatMessage]] = None,
+    messages: Optional[Sequence[TagChild | ChatMessageDict]] = None,
     placeholder: str = "Enter a message...",
     width: CssUnit = "min(680px, 100%)",
     height: CssUnit = "auto",

shiny/ui/_chat_normalize.py

@@ -4,7 +4,7 @@
 
 from htmltools import HTML, Tagifiable
 
-from ._chat_types import ChatUIMessage
+from ._chat_types import ChatMessage
 
 if TYPE_CHECKING:
     from anthropic.types import Message as AnthropicMessage
@@ -25,11 +25,11 @@ class GenerateContentResponse:
 
 class BaseMessageNormalizer(ABC):
     @abstractmethod
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         pass
 
     @abstractmethod
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         pass
 
     @abstractmethod
@@ -42,13 +42,13 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class StringNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast(Optional[str], message)
-        return ChatUIMessage(content=x or "", role="assistant")
+        return ChatMessage(content=x or "", role="assistant")
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast(Optional[str], chunk)
-        return ChatUIMessage(content=x or "", role="assistant")
+        return ChatMessage(content=x or "", role="assistant")
 
     def can_normalize(self, message: Any) -> bool:
         return isinstance(message, (str, HTML)) or message is None
@@ -58,17 +58,17 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class DictNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("dict[str, Any]", message)
         if "content" not in x:
             raise ValueError("Message must have 'content' key")
-        return ChatUIMessage(content=x["content"], role=x.get("role", "assistant"))
+        return ChatMessage(content=x["content"], role=x.get("role", "assistant"))
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("dict[str, Any]", chunk)
         if "content" not in x:
             raise ValueError("Message must have 'content' key")
-        return ChatUIMessage(content=x["content"], role=x.get("role", "assistant"))
+        return ChatMessage(content=x["content"], role=x.get("role", "assistant"))
 
     def can_normalize(self, message: Any) -> bool:
         return isinstance(message, dict)
@@ -78,11 +78,11 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class TagifiableNormalizer(DictNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("Tagifiable", message)
         return super().normalize({"content": x})
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("Tagifiable", chunk)
         return super().normalize_chunk({"content": x})
 
@@ -94,23 +94,23 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class LangChainNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("BaseMessage", message)
         if isinstance(x.content, list):  # type: ignore
             raise ValueError(
                 "The `message.content` provided seems to represent numerous messages. "
                 "Consider iterating over `message.content` and calling .append_message() on each iteration."
             )
-        return ChatUIMessage(content=x.content, role="assistant")
+        return ChatMessage(content=x.content, role="assistant")
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("BaseMessageChunk", chunk)
         if isinstance(x.content, list):  # type: ignore
             raise ValueError(
                 "The `message.content` provided seems to represent numerous messages. "
                 "Consider iterating over `message.content` and calling .append_message() on each iteration."
             )
-        return ChatUIMessage(content=x.content, role="assistant")
+        return ChatMessage(content=x.content, role="assistant")
 
     def can_normalize(self, message: Any) -> bool:
         try:
@@ -130,11 +130,11 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class OpenAINormalizer(StringNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("ChatCompletion", message)
         return super().normalize(x.choices[0].message.content)
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("ChatCompletionChunk", chunk)
         return super().normalize_chunk(x.choices[0].delta.content)
 
@@ -156,17 +156,17 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class AnthropicNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("AnthropicMessage", message)
         content = x.content[0]
         if content.type != "text":
             raise ValueError(
                 f"Anthropic message type {content.type} not supported. "
                 "Only 'text' type is currently supported"
             )
-        return ChatUIMessage(content=content.text, role="assistant")
+        return ChatMessage(content=content.text, role="assistant")
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("MessageStreamEvent", chunk)
         content = ""
         if x.type == "content_block_delta":
@@ -177,7 +177,7 @@ def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
                 )
             content = x.delta.text
 
-        return ChatUIMessage(content=content, role="assistant")
+        return ChatMessage(content=content, role="assistant")
 
     def can_normalize(self, message: Any) -> bool:
         try:
@@ -214,13 +214,13 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class GoogleNormalizer(BaseMessageNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("GenerateContentResponse", message)
-        return ChatUIMessage(content=x.text, role="assistant")
+        return ChatMessage(content=x.text, role="assistant")
 
-    def normalize_chunk(self, chunk: Any) -> ChatUIMessage:
+    def normalize_chunk(self, chunk: Any) -> ChatMessage:
         x = cast("GenerateContentResponse", chunk)
-        return ChatUIMessage(content=x.text, role="assistant")
+        return ChatMessage(content=x.text, role="assistant")
 
     def can_normalize(self, message: Any) -> bool:
         try:
@@ -238,11 +238,11 @@ def can_normalize_chunk(self, chunk: Any) -> bool:
 
 
 class OllamaNormalizer(DictNormalizer):
-    def normalize(self, message: Any) -> ChatUIMessage:
+    def normalize(self, message: Any) -> ChatMessage:
         x = cast("dict[str, Any]", message["message"])
         return super().normalize(x)
 
-    def normalize_chunk(self, chunk: "dict[str, Any]") -> ChatUIMessage:
+    def normalize_chunk(self, chunk: "dict[str, Any]") -> ChatMessage:
         msg = cast("dict[str, Any]", chunk["message"])
         return super().normalize_chunk(msg)
 
@@ -295,7 +295,7 @@ def register(
 message_normalizer_registry = NormalizerRegistry()
 
 
-def normalize_message(message: Any) -> ChatUIMessage:
+def normalize_message(message: Any) -> ChatMessage:
     strategies = message_normalizer_registry._strategies
     for strategy in strategies.values():
         if strategy.can_normalize(message):
@@ -306,7 +306,7 @@ def normalize_message(message: Any) -> ChatUIMessage:
     )
 
 
-def normalize_message_chunk(chunk: Any) -> ChatUIMessage:
+def normalize_message_chunk(chunk: Any) -> ChatMessage:
     strategies = message_normalizer_registry._strategies
     for strategy in strategies.values():
         if strategy.can_normalize_chunk(chunk):

shiny/ui/_chat_provider_types.py

@@ -1,7 +1,7 @@
 import sys
 from typing import TYPE_CHECKING, Literal, Union
 
-from ._chat_types import ChatMessage
+from ._chat_types import ChatMessageDict
 
 if TYPE_CHECKING:
     from anthropic.types import MessageParam as AnthropicMessage
@@ -47,7 +47,7 @@
 # TODO: use a strategy pattern to allow others to register
 # their own message formats
 def as_provider_message(
-    message: ChatMessage, format: ProviderMessageFormat
+    message: ChatMessageDict, format: ProviderMessageFormat
 ) -> "ProviderMessage":
     if format == "anthropic":
         return as_anthropic_message(message)
@@ -62,7 +62,7 @@ def as_provider_message(
     raise ValueError(f"Unknown format: {format}")
 
 
-def as_anthropic_message(message: ChatMessage) -> "AnthropicMessage":
+def as_anthropic_message(message: ChatMessageDict) -> "AnthropicMessage":
     from anthropic.types import MessageParam as AnthropicMessage
 
     if message["role"] == "system":
@@ -72,7 +72,7 @@ def as_anthropic_message(message: ChatMessage) -> "AnthropicMessage":
     return AnthropicMessage(content=message["content"], role=message["role"])
 
 
-def as_google_message(message: ChatMessage) -> "GoogleMessage":
+def as_google_message(message: ChatMessageDict) -> "GoogleMessage":
     if sys.version_info < (3, 9):
         raise ValueError("Google requires Python 3.9")
 
@@ -89,7 +89,7 @@ def as_google_message(message: ChatMessage) -> "GoogleMessage":
     return gtypes.ContentDict(parts=[message["content"]], role=role)
 
 
-def as_langchain_message(message: ChatMessage) -> "LangChainMessage":
+def as_langchain_message(message: ChatMessageDict) -> "LangChainMessage":
     from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
     content = message["content"]
@@ -103,7 +103,7 @@ def as_langchain_message(message: ChatMessage) -> "LangChainMessage":
     raise ValueError(f"Unknown role: {message['role']}")
 
 
-def as_openai_message(message: ChatMessage) -> "OpenAIMessage":
+def as_openai_message(message: ChatMessageDict) -> "OpenAIMessage":
     from openai.types.chat import (
         ChatCompletionAssistantMessageParam,
         ChatCompletionSystemMessageParam,
@@ -121,7 +121,7 @@ def as_openai_message(message: ChatMessage) -> "OpenAIMessage":
     raise ValueError(f"Unknown role: {role}")
 
 
-def as_ollama_message(message: ChatMessage) -> "OllamaMessage":
+def as_ollama_message(message: ChatMessageDict) -> "OllamaMessage":
     from ollama import Message as OllamaMessage
 
     return OllamaMessage(content=message["content"], role=message["role"])

shiny/ui/_chat_types.py

@@ -13,12 +13,12 @@
 
 # TODO: content should probably be [{"type": "text", "content": "..."}, {"type": "image", ...}]
 # in order to support multiple content types...
-class ChatMessage(TypedDict):
+class ChatMessageDict(TypedDict):
     content: str
     role: Role
 
 
-class ChatUIMessage:
+class ChatMessage:
     def __init__(
         self,
         content: TagChild,
@@ -50,7 +50,7 @@ class TransformedMessage:
     html_deps: list[dict[str, str]] | None = None
 
     @classmethod
-    def from_chat_message(cls, message: ChatUIMessage) -> "TransformedMessage":
+    def from_chat_message(cls, message: ChatMessage) -> "TransformedMessage":
         if message.role == "user":
             transform_key = "content_server"
             pre_transform_key = "content_client"