Update app.py

app.py

@@ -9,6 +9,8 @@ st.title("Llama Chatbot")
 # Get the API key from the environment variable
 API_KEY = os.getenv("HF_TOKEN")
 
+icons = {"assistant": "robot.png", "user": "man-kddi.png"}
+
 if API_KEY is None:
     st.error("API key not found. Please set the HF_TOKEN secret in your Hugging Face Space.")
     st.stop()
@@ -20,14 +22,6 @@ headers = {
     "Authorization": f"Bearer {API_KEY}"
 }
 
-def response_generator():
-    response = requests.post(API_URL, headers=headers, json={"inputs": json.dumps(st.session_state.messages)})
-    print("JSON:"+json.dumps(st.session_state.messages))
-    if response.status_code == 200:
-        yield response.json()[0]['generated_text']
-    else:
-        yield "Error: Unable to fetch response from the model."
-
 if "messages" not in st.session_state:
     st.session_state.messages = [{"role": "system", "content": 
                  """
@@ -37,21 +31,36 @@ if "messages" not in st.session_state:
                  """
                  }]
 
-# Display chat messages from history on app rerun
+def handle_query(query):
+    response = requests.post(API_URL, headers=headers, json={"inputs": json.dumps(st.session_state.messages)})
+    print("JSON:"+json.dumps(st.session_state.messages))
+    if response.status_code == 200:
+        return response.json()[0]['generated_text']
+    else:
+        return "Error: Unable to fetch response from the model."
+
+def streamer(text):
+    for i in text:
+        yield i
+        time.sleep(0.001)
+
+if 'messages' not in st.session_state:
+    st.session_state.messages = [{'role': 'assistant', "content": 'ask me anything'}]
+
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Accept user input
-if prompt := st.chat_input("What is up?"):
-    # 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})
-
-# Display assistant response in chat message container
-with st.chat_message("assistant"):
-    response = st.write_stream(response_generator())
-# Add assistant response to chat history
-st.session_state.messages.append({"role": "assistant", "content": response})
+    with st.chat_message(message['role'], avatar=icons[message['role']]):
+        st.write(message['content'])
+
+user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
+
+if user_prompt and (uploaded_file):
+    st.session_state.messages.append({'role': 'user', "content": user_prompt})
+    with st.chat_message("user", avatar="man-kddi.png"):
+        st.write(user_prompt)
+
+    # Trigger assistant's response retrieval and update UI
+    with st.spinner("Thinking..."):
+        response = handle_query(user_prompt)
+    with st.chat_message("user", avatar="robot.png"):
+        st.write_stream(streamer(response))
+    st.session_state.messages.append({'role': 'assistant', "content": response})