chat_to_pretty_markdown_simple.real_role update

h2oai · Sep 9, 2024 · 0f14fdd · 0f14fdd
1 parent 7a1a1d4
commit 0f14fdd
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/openai_server/chat_history_render.py b/openai_server/chat_history_render.py
@@ -122,13 +122,18 @@ def chat_to_pretty_markdown_simple(
             if role.lower() == "user"
             else "ℹ️"
         )
-        real_role = (
-            assistant_name
-            if role.lower() == "assistant"
-            else user_name
-            if role.lower() == "user"
-            else dummy_name
-        )
+        if 'name' in message:
+            # turns 'chat_agent' to 'Chat Agent'
+            real_role = message['name']
+            real_role = ' '.join(word.capitalize() for word in real_role.split('_'))
+        else:
+            real_role = (
+                assistant_name
+                if role.lower() == "assistant"
+                else user_name
+                if role.lower() == "user"
+                else dummy_name
+            )
 
         # Format the role
         if cute: