Update tests

Author: cpsievert

tests/pytest/test_chat.py

@@ -11,9 +11,8 @@
 from shiny.session import session_context
 from shiny.types import MISSING
 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
+from shiny.ui._chat_types import ChatMessage, ChatUIMessage
 
 # ----------------------------------------------------------------------
 # Helpers
@@ -52,31 +51,22 @@ def generate_content(token_count: int) -> str:
             return " ".join(["foo" for _ in range(1, n)])
 
         msgs = (
-            as_transformed_message(
-                {
-                    "content": generate_content(102),
-                    "role": "system",
-                }
-            ),
+            ChatUIMessage(
+                content=generate_content(102), role="system"
+            ).as_transformed_message(),
         )
 
         # Throws since system message is too long
         with pytest.raises(ValueError):
             chat._trim_messages(msgs, token_limits=(100, 0), format=MISSING)
 
         msgs = (
-            as_transformed_message(
-                {
-                    "content": generate_content(100),
-                    "role": "system",
-                }
-            ),
-            as_transformed_message(
-                {
-                    "content": generate_content(2),
-                    "role": "user",
-                }
-            ),
+            ChatUIMessage(
+                content=generate_content(100), role="system"
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=generate_content(2), role="user"
+            ).as_transformed_message(),
         )
 
         # Throws since only the system message fits
@@ -92,30 +82,24 @@ def generate_content(token_count: int) -> str:
         content3 = generate_content(2)
 
         msgs = (
-            as_transformed_message(
-                {
-                    "content": content1,
-                    "role": "system",
-                }
-            ),
-            as_transformed_message(
-                {
-                    "content": content2,
-                    "role": "user",
-                }
-            ),
-            as_transformed_message(
-                {
-                    "content": content3,
-                    "role": "user",
-                }
-            ),
+            ChatUIMessage(
+                content=content1,
+                role="system",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content2,
+                role="user",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content3,
+                role="user",
+            ).as_transformed_message(),
         )
 
         # Should discard the 1st user message
         trimmed = chat._trim_messages(msgs, token_limits=(103, 0), format=MISSING)
         assert len(trimmed) == 2
-        contents = [msg["content_server"] for msg in trimmed]
+        contents = [msg.content_server for msg in trimmed]
         assert contents == [content1, content3]
 
         content1 = generate_content(50)
@@ -124,38 +108,48 @@ def generate_content(token_count: int) -> str:
         content4 = generate_content(2)
 
         msgs = (
-            as_transformed_message(
-                {"content": content1, "role": "system"},
-            ),
-            as_transformed_message(
-                {"content": content2, "role": "user"},
-            ),
-            as_transformed_message(
-                {"content": content3, "role": "system"},
-            ),
-            as_transformed_message(
-                {"content": content4, "role": "user"},
-            ),
+            ChatUIMessage(
+                content=content1,
+                role="system",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content2,
+                role="user",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content3,
+                role="system",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content4,
+                role="user",
+            ).as_transformed_message(),
         )
 
         # Should discard the 1st user message
         trimmed = chat._trim_messages(msgs, token_limits=(103, 0), format=MISSING)
         assert len(trimmed) == 3
-        contents = [msg["content_server"] for msg in trimmed]
+        contents = [msg.content_server for msg in trimmed]
         assert contents == [content1, content3, content4]
 
         content1 = generate_content(50)
         content2 = generate_content(10)
 
         msgs = (
-            as_transformed_message({"content": content1, "role": "assistant"}),
-            as_transformed_message({"content": content2, "role": "user"}),
+            ChatUIMessage(
+                content=content1,
+                role="assistant",
+            ).as_transformed_message(),
+            ChatUIMessage(
+                content=content2,
+                role="user",
+            ).as_transformed_message(),
         )
 
         # Anthropic requires 1st message to be a user message
         trimmed = chat._trim_messages(msgs, token_limits=(30, 0), format="anthropic")
         assert len(trimmed) == 1
-        contents = [msg["content_server"] for msg in trimmed]
+        contents = [msg.content_server for msg in trimmed]
         assert contents == [content2]
 
 
@@ -172,13 +166,15 @@ def generate_content(token_count: int) -> str:
 
 
 def test_string_normalization():
-    msg = normalize_message_chunk("Hello world!")
-    assert msg == {"content": "Hello world!", "role": "assistant"}
+    m = normalize_message_chunk("Hello world!")
+    assert m.content == "Hello world!"
+    assert m.role == "assistant"
 
 
 def test_dict_normalization():
-    msg = normalize_message_chunk({"content": "Hello world!", "role": "assistant"})
-    assert msg == {"content": "Hello world!", "role": "assistant"}
+    m = normalize_message_chunk({"content": "Hello world!", "role": "assistant"})
+    assert m.content == "Hello world!"
+    assert m.role == "assistant"
 
 
 def test_langchain_normalization():
@@ -194,11 +190,15 @@ def test_langchain_normalization():
 
     # Mock & normalize return value of BaseChatModel.invoke()
     msg = BaseMessage(content="Hello world!", role="assistant", type="foo")
-    assert normalize_message(msg) == {"content": "Hello world!", "role": "assistant"}
+    m = normalize_message(msg)
+    assert m.content == "Hello world!"
+    assert m.role == "assistant"
 
     # Mock & normalize return value of BaseChatModel.stream()
     chunk = BaseMessageChunk(content="Hello ", type="foo")
-    assert normalize_message_chunk(chunk) == {"content": "Hello ", "role": "assistant"}
+    m = normalize_message_chunk(chunk)
+    assert m.content == "Hello "
+    assert m.role == "assistant"
 
 
 def test_google_normalization():
@@ -255,7 +255,9 @@ def test_anthropic_normalization():
         usage=Usage(input_tokens=0, output_tokens=0),
     )
 
-    assert normalize_message(msg) == {"content": "Hello world!", "role": "assistant"}
+    m = normalize_message(msg)
+    assert m.content == "Hello world!"
+    assert m.role == "assistant"
 
     # Mock return object from Anthropic().messages.create(stream=True)
     chunk = RawContentBlockDeltaEvent(
@@ -264,7 +266,9 @@ def test_anthropic_normalization():
         index=0,
     )
 
-    assert normalize_message_chunk(chunk) == {"content": "Hello ", "role": "assistant"}
+    m = normalize_message_chunk(chunk)
+    assert m.content == "Hello "
+    assert m.role == "assistant"
 
 
 def test_openai_normalization():
@@ -309,8 +313,9 @@ def test_openai_normalization():
         created=int(datetime.now().timestamp()),
     )
 
-    msg = normalize_message(completion)
-    assert msg == {"content": "Hello world!", "role": "assistant"}
+    m = normalize_message(completion)
+    assert m.content == "Hello world!"
+    assert m.role == "assistant"
 
     # Mock return object from OpenAI().chat.completions.create(stream=True)
     chunk = ChatCompletionChunk(
@@ -329,8 +334,9 @@ def test_openai_normalization():
         ],
     )
 
-    msg = normalize_message_chunk(chunk)
-    assert msg == {"content": "Hello ", "role": "assistant"}
+    m = normalize_message_chunk(chunk)
+    assert m.content == "Hello "
+    assert m.role == "assistant"
 
 
 def test_ollama_normalization():
@@ -343,8 +349,13 @@ def test_ollama_normalization():
     )
 
     msg_dict = {"content": "Hello world!", "role": "assistant"}
-    assert normalize_message(msg) == msg_dict
-    assert normalize_message_chunk(msg) == msg_dict
+    m = normalize_message(msg)
+    assert m.content == msg_dict["content"]
+    assert m.role == msg_dict["role"]
+
+    m = normalize_message_chunk(msg)
+    assert m.content == msg_dict["content"]
+    assert m.role == msg_dict["role"]
 
 
 # ------------------------------------------------------------------------------------
@@ -403,9 +414,7 @@ def test_as_google_message():
 
 
 def test_as_langchain_message():
-    from langchain_core.language_models.base import (
-        LanguageModelInput,
-    )
+    from langchain_core.language_models.base import LanguageModelInput
     from langchain_core.language_models.base import (
         Sequence as LangchainSequence,  # pyright: ignore[reportPrivateImportUsage]
     )