import streamlit as st
import gradio as gr
from streamlit_gradio import st_gradio

def chat_with_model(input_text):
    # Call the chat model and return the response
    response = chat_model.predict(input_text)
    return response

# Load the chat model using Gradio
chat_model = gr.Interface.load("models/CogwiseAI/CogwiseAI-chatwithMS")

# Create a Streamlit app
def main():
    st.title("Cogwise Intelligent Assistant")
    st.markdown("---")
    
    # Create a text input for user interaction
    input_text = st.text_input("You are talking to an AI, ask any question.")
    
    # Process user input and display the response
    if st.button("Ask"):
        with st.spinner("Thinking..."):
            response = chat_with_model(input_text)
        st.markdown("---")
        st.write(response)

    # Create the Gradio interface
    gr_interface = gr.Interface(fn=chat_with_model, inputs="text", outputs="text")
    
    # Render the Gradio interface using the st_gradio function
    st_gradio(gr_interface)

if __name__ == "__main__":
    main()