import streamlit as st
import ollama

def main():
    st.title("Llama Chatbot")
    
    user_input = st.text_input("You:", "Why is the sky blue?")

    if st.button("Ask"):
        with st.spinner("Thinking..."):
            response = ollama.chat(model='llama2', messages=[{
                'role': 'user',
                'content': user_input,
            }])
        st.success("Llama says:")
        st.write(response['message']['content'])

if __name__ == "__main__":
    main()