Add system message support to user function

app.py

@@ -21,9 +21,11 @@ def clear_chat(chat_history_state, chat_message):
     chat_message = ''
     return chat_history_state, chat_message
 
-def user(message, history):
+def user(message, history, system_message=None):
     print(f"User message: {message}")
     history = history or []
+    if system_message:  # Check if a system message is provided and should be added
+        history.append({"role": "system", "content": system_message})
     history.append({"role": "user", "content": message})
     return history
 
@@ -77,16 +79,21 @@ with gr.Blocks() as demo:
     top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.95)
     chat_history_state = gr.State([])
 
-    def update_chatbot(message, chat_history):
+    def update_chatbot(message, chat_history, system_message=BASE_SYSTEM_MESSAGE):
         print("Updating chatbot...")
+        # Check if we're starting a new conversation or continuing an existing one
         if not chat_history or (chat_history and chat_history[-1]["role"] != "user"):
+            # Include the system message only if starting a new conversation
+            chat_history = user(message, chat_history, system_message if not chat_history else None)
+        else:
+            # If continuing an existing conversation, don't add the system message again
             chat_history = user(message, chat_history)
-        chat_history, _ = chat(chat_history, BASE_SYSTEM_MESSAGE, max_tokens.value, temperature.value, top_p.value, 40, 1.1)
+        chat_history, _ = chat(chat_history, system_message, max_tokens.value, temperature.value, top_p.value, 40, 1.1)
         # Format the chat history for the gr.Chatbot component
         formatted_chat_history = []
-        for user_msg, assistant_msg in zip([msg["content"] for msg in chat_history if msg["role"] == "user"],
-                                        [msg["content"] for msg in chat_history if msg["role"] == "assistant"]):
-            formatted_chat_history.append([user_msg, assistant_msg])
+        for entry in chat_history:
+            if entry["role"] in ["user", "assistant"]:
+                formatted_chat_history.append([entry["content"] if entry["role"] == "user" else "", entry["content"] if entry["role"] == "assistant" else ""])
         return formatted_chat_history, chat_history, ""