Update app.py

app.py

@@ -8,10 +8,10 @@ load_dotenv()
 # Initialize the client
 client = OpenAI(
     base_url="https://api-inference.huggingface.co/v1",
-    api_key=os.getenv('HUGGINGFACEHUB_API_TOKEN')  # Ensure the environment variable is set correctly
+    api_key=os.getenv('HUGGINGFACEHUB_API_TOKEN')
 )
 
-# Model link
+# Model link for Mistral
 model_link = "mistralai/Mistral-7B-Instruct-v0.2"
 
 def reset_conversation():
@@ -21,27 +21,27 @@ def reset_conversation():
     st.session_state.messages = []
     return None
 
+# Define the available models
+model_links = {
+    "Mistral-7B": "mistralai/Mistral-7B-Instruct-v0.2"
+}
+
 # Set the temperature value directly in the code
 temperature = 0.5
 
-# Add a button to clear conversation
-st.button('Reset Chat', on_click=reset_conversation)  # Reset button
+# Add reset button to clear conversation
+st.button('Reset Chat', on_click=reset_conversation)
+
+# Define a system message for the Mistral model
+system_message = {
+    "role": "system",
+    "content": "As a helpful and friendly assistant, provide concise and accurate responses to the user's queries."
+}
 
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-st.title("Mistral-7B Chatbot")
-
-# Function to prepare messages for the API
-def get_messages_for_api():
-    messages = []
-    if len(st.session_state.messages) % 2 == 1:  # If the last message is from the user
-        messages.append({"role": "assistant", "content": "Your assistant is ready."})  # Ensure the assistant starts or resumes
-    for message in st.session_state.messages:
-        messages.append(message)
-    return messages
-
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
@@ -49,34 +49,29 @@ 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)
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Prepare messages for the conversation context
-    all_messages = get_messages_for_api()
-
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
         try:
             response = client.chat.completions.create(
                 model=model_link,
-                messages=all_messages,
+                messages=[
+                    system_message,
+                    {"role": m["role"], "content": m["content"]}
+                    for m in st.session_state.messages
+                ],
                 temperature=temperature,
-                max_tokens=150  # Adjust the token limit as needed
+                max_tokens=3000,
             )
-
-            if 'choices' in response and response['choices']:
-                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.")
-        
+            st.markdown(response["choices"][0]["message"]["content"])
         except Exception as e:
-            st.markdown("An error occurred. Please try again later.")
-            st.markdown("Error details:")
-            st.markdown(str(e))
-            st.session_state.messages.append({"role": "assistant", "content": "An error occurred. Please try again later."})
+            st.write("An error occurred. Please try again later.")
+            st.write(e)
+
+    st.session_state.messages.append({"role": "assistant", "content": response["choices"][0]["message"]["content"]})