Chat.messages() now keeps system messages (#1509)

Author: cpsievert

shiny/ui/_chat.py

@@ -386,7 +386,9 @@ def messages(
             A tuple of chat messages.
         """
 
-        messages = self._get_trimmed_messages(token_limits=token_limits)
+        messages = self._messages()
+        if token_limits is not None:
+            messages = self._trim_messages(messages, token_limits)
 
         res: list[ChatMessage] = []
         for i, m in enumerate(messages):
@@ -741,34 +743,63 @@ def _store_message(
 
         return msg
 
-    def _get_trimmed_messages(
-        self,
-        *,
-        token_limits: tuple[int, int] | None = (4096, 1000),
+    @staticmethod
+    def _trim_messages(
+        messages: tuple[StoredMessage, ...],
+        token_limits: tuple[int, int] = (4096, 1000),
     ) -> tuple[StoredMessage, ...]:
-        messages = self._messages()
 
-        if token_limits is None:
-            return messages
+        n_total, n_reserve = token_limits
+        if n_total <= n_reserve:
+            raise ValueError(
+                f"Invalid token limits: {token_limits}. The 1st value must be greater "
+                "than the 2nd value."
+            )
 
-        # Can't trim if we don't have token counts
-        token_counts = [m["token_count"] for m in messages]
-        if None in token_counts:
-            return messages
+        # Since don't trim system messages, 1st obtain their total token count
+        # (so we can determine how many non-system messages can fit)
+        n_system_tokens: int = 0
+        n_system_messages: int = 0
+        n_other_messages: int = 0
+        for m in messages:
+            count = m["token_count"]
+            # Count can be None if the tokenizer is None
+            if count is None:
+                return messages
+            if m["role"] == "system":
+                n_system_tokens += count
+                n_system_messages += 1
+            else:
+                n_other_messages += 1
+
+        remaining_non_system_tokens = n_total - n_reserve - n_system_tokens
 
-        token_counts = cast("list[int]", token_counts)
+        if remaining_non_system_tokens <= 0:
+            raise ValueError(
+                f"System messages exceed `.messages(token_limits={token_limits})`. "
+                "Consider increasing the 1st value of `token_limit` or setting it to "
+                "`token_limit=None` to disable token limits."
+            )
 
-        # Take the newest messages up to the token limit
-        limit, reserve = token_limits
-        max_tokens = limit - reserve
         messages2: list[StoredMessage] = []
-        for i, m in enumerate(reversed(messages)):
-            if sum(token_counts[-i - 1 :]) > max_tokens:
-                break
-            messages2.append(m)
+        for m in reversed(messages):
+            if m["role"] == "system":
+                messages2.append(m)
+                continue
+            count = cast(int, m["token_count"])  # Already checked this
+            remaining_non_system_tokens -= count
+            if remaining_non_system_tokens >= 0:
+                messages2.append(m)
 
         messages2.reverse()
 
+        if len(messages2) == n_system_messages and n_other_messages > 0:
+            raise ValueError(
+                f"Only system messages fit within `.messages(token_limits={token_limits})`. "
+                "Consider increasing the 1st value of `token_limit` or setting it to "
+                "`token_limit=None` to disable token limits."
+            )
+
         return tuple(messages2)
 
     def user_input(self, transform: bool = False) -> str | None:

tests/pytest/test_chat.py

@@ -2,10 +2,127 @@
 
 import sys
 from datetime import datetime
+from typing import cast
 
+import pytest
+
+from shiny import Session
+from shiny._namespaces import ResolvedId, Root
+from shiny.session import session_context
+from shiny.ui import Chat
+from shiny.ui._chat import as_transformed_message
 from shiny.ui._chat_normalize import normalize_message, normalize_message_chunk
+from shiny.ui._chat_types import ChatMessage, StoredMessage
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+class _MockSession:
+    ns: ResolvedId = Root
+    app: object = None
+    id: str = "mock-session"
+
+    def on_ended(self, callback: object) -> None:
+        pass
+
+    def _increment_busy_count(self) -> None:
+        pass
+
+
+test_session = cast(Session, _MockSession())
+
+
+def as_stored_message(message: ChatMessage, token_count: int) -> StoredMessage:
+    msg = as_transformed_message(message)
+    return StoredMessage(
+        **msg,
+        token_count=token_count,
+    )
+
 
-# TODO: Feed these messages into an actual Chat() instance?
+# ----------------------------------------------------------------------
+# Unit tests for Chat._get_trimmed_messages()
+# ----------------------------------------------------------------------
+
+
+def test_chat_message_trimming():
+    with session_context(test_session):
+        chat = Chat(id="chat")
+
+        msgs = (
+            as_stored_message(
+                {"content": "System message", "role": "system"}, token_count=101
+            ),
+        )
+
+        # Throws since system message is too long
+        with pytest.raises(ValueError):
+            chat._trim_messages(msgs, token_limits=(100, 0))
+
+        msgs = (
+            as_stored_message(
+                {"content": "System message", "role": "system"}, token_count=100
+            ),
+            as_stored_message(
+                {"content": "User message", "role": "user"}, token_count=1
+            ),
+        )
+
+        # Throws since only the system message fits
+        with pytest.raises(ValueError):
+            chat._trim_messages(msgs, token_limits=(100, 0))
+
+        # Raising the limit should allow both messages to fit
+        trimmed = chat._trim_messages(msgs, token_limits=(102, 0))
+        assert len(trimmed) == 2
+        contents = [msg["content_server"] for msg in trimmed]
+        assert contents == ["System message", "User message"]
+
+        msgs = (
+            as_stored_message(
+                {"content": "System message", "role": "system"}, token_count=100
+            ),
+            as_stored_message(
+                {"content": "User message", "role": "user"}, token_count=10
+            ),
+            as_stored_message(
+                {"content": "User message 2", "role": "user"}, token_count=1
+            ),
+        )
+
+        # Should discard the 1st user message
+        trimmed = chat._trim_messages(msgs, token_limits=(102, 0))
+        assert len(trimmed) == 2
+        contents = [msg["content_server"] for msg in trimmed]
+        assert contents == ["System message", "User message 2"]
+
+        msgs = (
+            as_stored_message(
+                {"content": "System message", "role": "system"}, token_count=50
+            ),
+            as_stored_message(
+                {"content": "User message", "role": "user"}, token_count=10
+            ),
+            as_stored_message(
+                {"content": "System message 2", "role": "system"}, token_count=50
+            ),
+            as_stored_message(
+                {"content": "User message 2", "role": "user"}, token_count=1
+            ),
+        )
+
+        # Should discard the 1st user message
+        trimmed = chat._trim_messages(msgs, token_limits=(102, 0))
+        assert len(trimmed) == 3
+        contents = [msg["content_server"] for msg in trimmed]
+        assert contents == ["System message", "System message 2", "User message 2"]
+
+
+# ----------------------------------------------------------------------
+# Unit tests for normalize_message() and normalize_message_chunk()
+# ----------------------------------------------------------------------
 
 
 def test_string_normalization():