refactor: make ChatMessage its own component

widgetti · Nov 10, 2023 · 4153745 · 4153745
1 parent 1d70dc9
commit 4153745
Showing 1 changed file with 47 additions and 45 deletions.
diff --git a/wanderlust.py b/wanderlust.py
@@ -122,6 +122,51 @@ def Map():
     )
 
 
+@solara.component
+def ChatMessage(message):
+    # Catch "messages" that are actually tool calls
+    if isinstance(message, dict):
+        icon = "mdi-map" if message["output"] == "Map updated" else "mdi-map-marker"
+        solara.v.Icon(children=[icon], style_="padding-top: 10px;")
+        solara.Markdown(message["output"])
+    elif message.role == "user":
+        solara.Text(
+            message.content[0].text.value,
+            classes=["chat-message", "user-message"],
+        )
+    elif message.role == "assistant":
+        if message.content[0].text.value:
+            solara.v.Icon(
+                children=["mdi-compass-outline"],
+                style_="padding-top: 10px;",
+            )
+            solara.Markdown(message.content[0].text.value)
+        elif message.content.tool_calls:
+            solara.v.Icon(
+                children=["mdi-map"],
+                style_="padding-top: 10px;",
+            )
+            solara.Markdown("*Calling map functions*")
+        else:
+            solara.v.Icon(
+                children=["mdi-compass-outline"],
+                style_="padding-top: 10px;",
+            )
+            solara.Preformatted(
+                repr(message),
+                classes=["chat-message", "assistant-message"],
+            )
+    else:
+        solara.v.Icon(
+            children=["mdi-compass-outline"],
+            style_="padding-top: 10px;",
+        )
+        solara.Preformatted(
+            repr(message),
+            classes=["chat-message", "assistant-message"],
+        )
+
+
 @solara.component
 def ChatInterface():
     prompt = solara.use_reactive("")
@@ -196,51 +241,8 @@ def poll():
             ):
                 for message in reversed(messages.value):
                     with solara.Row(style={"align-items": "flex-start"}):
-                        # Catch "messages" that are actually tool calls
-                        if isinstance(message, dict):
-                            icon = (
-                                "mdi-map"
-                                if message["output"] == "Map updated"
-                                else "mdi-map-marker"
-                            )
-                            solara.v.Icon(children=[icon], style_="padding-top: 10px;")
-                            solara.Markdown(message["output"])
-                        elif message.role == "user":
-                            solara.Text(
-                                message.content[0].text.value,
-                                classes=["chat-message", "user-message"],
-                            )
-                        elif message.role == "assistant":
-                            if message.content[0].text.value:
-                                solara.v.Icon(
-                                    children=["mdi-compass-outline"],
-                                    style_="padding-top: 10px;",
-                                )
-                                solara.Markdown(message.content[0].text.value)
-                            elif message.content.tool_calls:
-                                solara.v.Icon(
-                                    children=["mdi-map"],
-                                    style_="padding-top: 10px;",
-                                )
-                                solara.Markdown("*Calling map functions*")
-                            else:
-                                solara.v.Icon(
-                                    children=["mdi-compass-outline"],
-                                    style_="padding-top: 10px;",
-                                )
-                                solara.Preformatted(
-                                    repr(message),
-                                    classes=["chat-message", "assistant-message"],
-                                )
-                        else:
-                            solara.v.Icon(
-                                children=["mdi-compass-outline"],
-                                style_="padding-top: 10px;",
-                            )
-                            solara.Preformatted(
-                                repr(message),
-                                classes=["chat-message", "assistant-message"],
-                            )
+                        ChatMessage(message)
+
         with solara.Column():
             solara.InputText(
                 label="Ask your question here",