Show chat model in assistant's reply, not user's

paulovcmedeiros · Mar 8, 2024 · e4f9e64 · e4f9e64
1 parent 138814d
commit e4f9e64
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/pyrobbot/app/app_page_templates.py b/pyrobbot/app/app_page_templates.py
@@ -259,7 +259,10 @@ def render_chat_history(self):
                 continue
             with st.chat_message(role, avatar=self.avatars.get(role)):
                 with contextlib.suppress(KeyError):
-                    st.caption(f"{message['chat_model']}, {message['timestamp']}")
+                    if role == "assistant":
+                        st.caption(message["chat_model"])
+                    else:
+                        st.caption(message["timestamp"])
                 st.markdown(message["content"])
                 with contextlib.suppress(KeyError):
                     if audio := message.get("reply_audio_file_path"):
@@ -403,7 +406,6 @@ def _render_chatbot_page(self):  # noqa: PLR0915
                             "name": self.chat_obj.username,
                             "content": prompt,
                             "timestamp": time_now,
-                            "chat_model": self.chat_obj.model,
                         }
                     )
 
@@ -418,6 +420,7 @@ def _render_chatbot_page(self):  # noqa: PLR0915
                                 "name": self.chat_obj.assistant_name,
                                 "content": reply["text"],
                                 "reply_audio_file_path": reply["audio"],
+                                "chat_model": self.chat_obj.model,
                             }
                         )