Update app.py

app.py

@@ -8,7 +8,7 @@ load_dotenv()
 # Initialize the client
 client = OpenAI(
     base_url="https://api-inference.huggingface.co/v1",
-    api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')  # Replace with your token
+    api_key=os.getenv('HUGGINGFACEHUB_API_TOKEN')  # Ensure the environment variable is set correctly
 )
 
 # Model link
@@ -35,15 +35,18 @@ st.title("Mistral-7B Chatbot")
 
 # Function to get streamed response
 def get_streamed_response(message, history):
-    all_message = [{
+    all_messages = []
+    # Include system prompt
+    all_messages.append({
         "role": "system",
-        "content": "From now on, you are an AI assistant knowledgeable in general topics. You can respond with relevant information and provide concise, friendly replies. Always maintain a helpful and neutral tone. Ensure to be concise to keep the conversation flowing smoothly."
-    }]
-    
+        "content": "You are a helpful assistant knowledgeable in general topics. Respond concisely and maintain a friendly and neutral tone."
+    })
+    # Add previous messages to history
     for msg in history:
-        all_message.append(msg)
-    all_message.append({"role": "user", "content": message})
-    return all_message
+        all_messages.append(msg)
+    # Add the latest user message
+    all_messages.append({"role": "user", "content": message})
+    return all_messages
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
@@ -52,7 +55,6 @@ for message in st.session_state.messages:
 
 # Accept user input
 if prompt := st.chat_input("Type your message here..."):
-
     # Display user message in chat message container
     with st.chat_message("user"):
         st.markdown(prompt)
@@ -65,17 +67,18 @@ if prompt := st.chat_input("Type your message here..."):
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
         try:
+            all_messages = get_streamed_response(prompt, history)
             response = client.chat.completions.create(
                 model=model_link,
-                messages=get_streamed_response(prompt, history),
+                messages=all_messages,
                 temperature=temperature,
                 max_tokens=150  # Adjust the token limit as needed
             )
 
-            # Ensure response handling is safe
             if 'choices' in response and response['choices']:
-                st.markdown(response['choices'][0]['message']['content'])
-                st.session_state.messages.append({"role": "assistant", "content": response['choices'][0]['message']['content']})
+                content = response['choices'][0]['message']['content']
+                st.markdown(content)
+                st.session_state.messages.append({"role": "assistant", "content": content})
             else:
                 st.markdown("No response received from the assistant.")