Update app.py

app.py

@@ -16,28 +16,23 @@ headers = {
 # Base system message
 BASE_SYSTEM_MESSAGE = "I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning."
 
-
 def clear_chat():
     """Clears the chat history and message state."""
     print("Clearing chat...")
     chat_history_state.value = []
     chatbot.textbox.value = ""
 
-
-def user(message, history, system_message=None):
+def user(message, history):
     """Updates the chat history with the user message."""
     print(f"User message: {message}")
     history = history or []
-    if system_message:
-        history.append({"role": "system", "content": system_message})
     history.append({"role": "user", "content": message})
     return history
 
-
 def call_nvidia_api(history, system_message, max_tokens, temperature, top_p):
     """Calls the NVIDIA API to generate a response."""
     messages = [{"role": "system", "content": system_message}]
-    messages.extend([{"role": "user", "content": h[0]} for h in history])
+    messages.extend([{"role": role, "content": content} for role, content in history])
 
     payload = {
         "messages": messages,
@@ -58,7 +53,7 @@ def call_nvidia_api(history, system_message, max_tokens, temperature, top_p):
     print(f"Payload recebido: {response_body}")
     if response_body.get("choices"):
         assistant_message = response_body["choices"][0]["message"]["content"]
-        history.append(["", assistant_message])
+        history.append({"role": "assistant", "content": assistant_message})
     return history
 
 def chatbot_submit(message, chat_history, system_message, max_tokens_val, temperature_val, top_p_val):
@@ -66,19 +61,19 @@ def chatbot_submit(message, chat_history, system_message, max_tokens_val, temper
     print("Updating chatbot...")
 
     # Adiciona a mensagem do usuário ao histórico
-    chat_history.append([message, ""])
+    chat_history = user(message, chat_history)
 
     # Chama a API da NVIDIA para gerar uma resposta
     chat_history = call_nvidia_api(chat_history, system_message, max_tokens_val, temperature_val, top_p_val)
 
     # Extrai apenas a mensagem do assistente da resposta
-    if chat_history and chat_history[-1][1]:
-        assistant_message = chat_history[-1][1]
+    if chat_history and chat_history[-1]["role"] == "assistant":
+        assistant_message = chat_history[-1]["content"]
     else:
         assistant_message = "Desculpe, ocorreu um erro ao gerar a resposta."
 
     return assistant_message, chat_history
-    
+
 chat_history_state = gr.State([])
 system_msg = gr.Textbox(BASE_SYSTEM_MESSAGE, label="System Message", placeholder="System prompt.", lines=5)
 max_tokens = gr.Slider(20, 1024, label="Max Tokens", step=20, value=1024)