Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,6 +2,7 @@ import os
 os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"
 from langchain_huggingface import HuggingFaceEndpoint
 import streamlit as st
+from langchain_core.messages import HumanMessage
 
 # constants
 QUESTION = "Compute the integral of f(x) = x^2."
@@ -116,7 +117,7 @@ with st.container():
                             f"{QUESTION}\n"
                             "[/INST]"
                         )
-                        st.session_state.response = llm.invoke(full_prompt)
+                        st.session_state.response = llm.invoke([HumanMessage(content=full_prompt)])
                         st.session_state.retry_count = 0
                         st.session_state.full_prompt = full_prompt  # Save prompt for retry
                         st.session_state.button_clicked = "Explain the question"
@@ -131,7 +132,7 @@ with st.container():
                                 f"{QUESTION}\n"
                                 "[/INST]"
                             )
-                            st.session_state.response = llm.invoke(full_prompt)
+                            st.session_state.response = llm.invoke([HumanMessage(content=full_prompt)])
                             st.session_state.retry_count = 0
                             st.session_state.full_prompt = full_prompt  # Save prompt for retry
                             st.session_state.button_clicked = "Give an example"
@@ -146,7 +147,7 @@ with st.container():
                                 f"{QUESTION}\n"
                                 "[/INST]"
                             )
-                            st.session_state.response = llm.invoke(full_prompt)
+                            st.session_state.response = llm.invoke([HumanMessage(content=full_prompt)])
                             st.session_state.retry_count = 0
                             st.session_state.full_prompt = full_prompt  # Save prompt for retry
                             st.session_state.button_clicked = "Who cares?"
@@ -166,6 +167,8 @@ with st.container():
                             st.session_state.response = alt_llm.invoke(
                                 st.session_state.full_prompt
                             )
+                        st.info(st.session_state.response)
+                        
 # Optional: Add footer or spacing
 st.markdown("<br><br>", unsafe_allow_html=True)