Move away from inheriting from TypedDict for internal Chat classes

Author: cpsievert

shiny/session/_utils.py

@@ -10,8 +10,6 @@
 from contextvars import ContextVar, Token
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 
-from htmltools import TagChild
-
 if TYPE_CHECKING:
     from ._session import Session
 
@@ -134,17 +132,6 @@ def require_active_session(session: Optional[Session]) -> Session:
     return session
 
 
-def process_ui(ui: TagChild) -> tuple[str, list[dict[str, str]]]:
-    """
-    Process a UI element with the session, returning the HTML and dependencies.
-    """
-    if isinstance(ui, (str, float, int)):
-        return str(ui), []
-    session = require_active_session(None)
-    res = session._process_ui(ui)
-    return res["html"], res["deps"]
-
-
 # Ideally I'd love not to limit the types for T, but if I don't, the type checker has
 # trouble figuring out what `T` is supposed to be when run_thunk is actually used. For
 # now, just keep expanding the possible types, as needed.

shiny/ui/_chat.py

@@ -38,7 +38,7 @@
     as_provider_message,
 )
 from ._chat_tokenizer import TokenEncoding, TokenizersEncoding, get_default_tokenizer
-from ._chat_types import ChatMessage, ClientMessage, TransformedMessage
+from ._chat_types import ChatMessage, ChatUIMessage, ClientMessage, TransformedMessage
 from ._html_deps_py_shiny import chat_deps
 from .fill import as_fill_item, as_fillable_container
 
@@ -240,7 +240,7 @@ async def _init_chat():
             @reactive.effect(priority=9999)
             @reactive.event(self._user_input)
             async def _on_user_input():
-                msg = ChatMessage(content=self._user_input(), role="user")
+                msg = ChatUIMessage(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())
@@ -251,7 +251,7 @@ 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(as_transformed_message(msg), index=n_pre)
+                    self._store_message(msg.as_transformed_message(), index=n_pre)
                     await self._remove_loading_message()
                     self._suspend_input_handler = True
 
@@ -492,14 +492,17 @@ def messages(
         res: list[ChatMessage | ProviderMessage] = []
         for i, m in enumerate(messages):
             transform = False
-            if m["role"] == "assistant":
+            if m.role == "assistant":
                 transform = transform_assistant
-            elif m["role"] == "user":
+            elif m.role == "user":
                 transform = transform_user == "all" or (
                     transform_user == "last" and i == len(messages) - 1
                 )
-            content_key = m["transform_key" if transform else "pre_transform_key"]
-            chat_msg = ChatMessage(content=str(m[content_key]), role=m["role"])
+            content_key = getattr(
+                m, "transform_key" if transform else "pre_transform_key"
+            )
+            content = getattr(m, content_key)
+            chat_msg = ChatMessage(content=str(content), role=m.role)
             if not isinstance(format, MISSING_TYPE):
                 chat_msg = as_provider_message(chat_msg, format)
             res.append(chat_msg)
@@ -593,9 +596,9 @@ async def _append_message(
         else:
             msg = normalize_message_chunk(message)
             # Update the current stream message
-            chunk_content = msg["content"]
+            chunk_content = msg.content
             self._current_stream_message += chunk_content
-            msg["content"] = self._current_stream_message
+            msg.content = self._current_stream_message
             if chunk == "end":
                 self._current_stream_message = ""
 
@@ -771,7 +774,7 @@ async def _send_append_message(
         chunk: ChunkOption = False,
         icon: HTML | Tag | TagList | None = None,
     ):
-        if message["role"] == "system":
+        if message.role == "system":
             # System messages are not displayed in the UI
             return
 
@@ -786,21 +789,21 @@ async def _send_append_message(
         elif chunk == "end":
             chunk_type = "message_end"
 
-        content = message["content_client"]
+        content = message.content_client
         content_type = "html" if isinstance(content, HTML) else "markdown"
 
         # TODO: pass along dependencies for both content and icon (if any)
         msg = ClientMessage(
             content=str(content),
-            role=message["role"],
+            role=message.role,
             content_type=content_type,
             chunk_type=chunk_type,
         )
 
         if icon is not None:
             msg["icon"] = str(icon)
 
-        deps = message.get("html_deps", [])
+        deps = message.html_deps
         if deps:
             msg["html_deps"] = deps
 
@@ -928,19 +931,19 @@ async def _transform_wrapper(content: str, chunk: str, done: bool):
 
     async def _transform_message(
         self,
-        message: ChatMessage,
+        message: ChatUIMessage,
         chunk: ChunkOption = False,
         chunk_content: str | None = None,
     ) -> TransformedMessage | None:
-        res = as_transformed_message(message)
-        key = res["transform_key"]
+        res = message.as_transformed_message()
+        key = res.transform_key
 
-        if message["role"] == "user" and self._transform_user is not None:
-            content = await self._transform_user(message["content"])
+        if message.role == "user" and self._transform_user is not None:
+            content = await self._transform_user(message.content)
 
-        elif message["role"] == "assistant" and self._transform_assistant is not None:
+        elif message.role == "assistant" and self._transform_assistant is not None:
             content = await self._transform_assistant(
-                message["content"],
+                message.content,
                 chunk_content or "",
                 chunk == "end" or chunk is False,
             )
@@ -975,7 +978,7 @@ def _store_message(
         messages.insert(index, message)
 
         self._messages.set(tuple(messages))
-        if message["role"] == "user":
+        if message.role == "user":
             self._latest_user_input.set(message)
 
         return None
@@ -1000,9 +1003,9 @@ def _trim_messages(
         n_other_messages: int = 0
         token_counts: list[int] = []
         for m in messages:
-            count = self._get_token_count(m["content_server"])
+            count = self._get_token_count(m.content_server)
             token_counts.append(count)
-            if m["role"] == "system":
+            if m.role == "system":
                 n_system_tokens += count
                 n_system_messages += 1
             else:
@@ -1023,7 +1026,7 @@ def _trim_messages(
         n_other_messages2: int = 0
         token_counts.reverse()
         for i, m in enumerate(reversed(messages)):
-            if m["role"] == "system":
+            if m.role == "system":
                 messages2.append(m)
                 continue
             remaining_non_system_tokens -= token_counts[i]
@@ -1046,13 +1049,13 @@ def _trim_anthropic_messages(
         self,
         messages: tuple[TransformedMessage, ...],
     ) -> tuple[TransformedMessage, ...]:
-        if any(m["role"] == "system" for m in messages):
+        if any(m.role == "system" for m in messages):
             raise ValueError(
                 "Anthropic requires a system prompt to be specified in it's `.create()` method "
                 "(not in the chat messages with `role: system`)."
             )
         for i, m in enumerate(messages):
-            if m["role"] == "user":
+            if m.role == "user":
                 return messages[i:]
 
         return ()
@@ -1098,7 +1101,8 @@ def user_input(self, transform: bool = False) -> str | None:
         if msg is None:
             return None
         key = "content_server" if transform else "content_client"
-        return str(msg[key])
+        val = getattr(msg, key)
+        return str(val)
 
     def _user_input(self) -> str:
         id = self.user_input_id
@@ -1361,27 +1365,4 @@ def chat_ui(
     return res
 
 
-def as_transformed_message(message: ChatMessage) -> TransformedMessage:
-    if message["role"] == "user":
-        transform_key = "content_server"
-        pre_transform_key = "content_client"
-    else:
-        transform_key = "content_client"
-        pre_transform_key = "content_server"
-
-    res = TransformedMessage(
-        content_client=message["content"],
-        content_server=message["content"],
-        role=message["role"],
-        transform_key=transform_key,
-        pre_transform_key=pre_transform_key,
-    )
-
-    deps = message.get("html_deps", [])
-    if deps:
-        res["html_deps"] = deps
-
-    return res
-
-
 CHAT_INSTANCES: WeakValueDictionary[str, Chat] = WeakValueDictionary()