Update app.py

app.py

@@ -2,35 +2,60 @@ import streamlit as st
 from langchain.chat_models import HuggingFaceHub
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
-# StreamLit UI
+# ------------------------
+# Streamlit UI
+# ------------------------
 st.set_page_config(page_title="LangChain Demo", page_icon=":robot:")
 st.header("MHRV Chatbot")
 
 if "sessionMessages" not in st.session_state:
     st.session_state.sessionMessages = [
-        SystemMessage(content="You are customer support chatbot from a website.")
+        SystemMessage(content="You are a customer support chatbot for a website.")
     ]
 
-# Load HuggingFace model
+# ------------------------
+# Load Hugging Face Model
+# ------------------------
 chat = HuggingFaceHub(
     repo_id="OpenAssistant/oasst-sft-4-pythia-12b",
     model_kwargs={"temperature": 0, "max_new_tokens": 512}
 )
 
+# ------------------------
+# Helper Functions
+# ------------------------
 def load_answer(question):
+    # Add human message
     st.session_state.sessionMessages.append(HumanMessage(content=question))
-    assistant_answer = chat(st.session_state.sessionMessages)
+
+    # Convert messages to a single prompt string
+    prompt = ""
+    for msg in st.session_state.sessionMessages:
+        if isinstance(msg, SystemMessage):
+            prompt += f"System: {msg.content}\n"
+        elif isinstance(msg, HumanMessage):
+            prompt += f"Human: {msg.content}\n"
+        elif isinstance(msg, AIMessage):
+            prompt += f"AI: {msg.content}\n"
+
+    # Get AI response
+    assistant_answer = chat(prompt)
+
+    # Add AI message to session
     st.session_state.sessionMessages.append(AIMessage(content=assistant_answer.content))
+
     return assistant_answer.content
 
 def get_text():
-    input_text = st.text_input("You: ", key="input")
-    return input_text
+    return st.text_input("You: ", key="input")
 
+# ------------------------
+# Main App
+# ------------------------
 user_input = get_text()
-submit = st.button('Generate')  
+submit = st.button("Generate")
 
 if submit and user_input:
     response = load_answer(user_input)
     st.subheader("Answer:")
-    st.write(response, key=1)
+    st.write(response)