import os
import streamlit as st
from autogen import AssistantAgent, UserProxyAgent

# Configure Groq API
llm_config = {
    "config_list": [
        {
            "model": "llama3-8b-8192",
            "api_key": os.environ.get("GROQ_API_KEY"),
            "base_url": "https://api.groq.com/openai/v1"
        }
    ]
}

# Initialize AutoGen agents
assistant = AssistantAgent("assistant", llm_config=llm_config)
user_proxy = UserProxyAgent("user_proxy", code_execution_config=False)

def main():
    st.title("Chatbot Interface")
    st.write("Ask me anything!")

    # Initialize session state for chat history
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    # Function to handle chat interaction
    def respond(message):
        if not message.strip():
            return "Error: No input received."

        try:
            user_proxy.initiate_chat(assistant, message=message)
            response = assistant.last_message().get("content", "")
            st.session_state.chat_history.append({"role": "user", "content": message})
            st.session_state.chat_history.append({"role": "assistant", "content": response})
        except Exception as e:
            error_message = f"Error: {str(e)}"
            st.session_state.chat_history.append({"role": "assistant", "content": error_message})

    # Display chat history
    for chat in st.session_state.chat_history:
        if chat["role"] == "user":
            st.markdown(f"**You:** {chat['content']}")
        else:
            st.markdown(f"**Assistant:** {chat['content']}")

    # User input form
    user_input = st.text_input("Type your message here...", "")
    if st.button('Send'):
        respond(user_input)

if __name__ == "__main__":
    main()