app.py

import streamlit as st
import openai
from typing import List, Dict

# Set up page configuration
st.set_page_config(
    page_title="Grok AI Assistant",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for enhanced styling
st.markdown("""
<style>
.main-container {
    background-color: #0E1117;
    color: #FFFFFF;
    font-family: 'Roboto', sans-serif;
}
.stTextInput>div>div>input {
    background-color: #262730;
    color: white;
    border: 2px solid #4A4A4A;
}
.stButton>button {
    background-color: #FF4B4B;
    color: white;
    border: none;
    padding: 10px 20px;
    border-radius: 5px;
    transition: all 0.3s ease;
}
.stButton>button:hover {
    background-color: #FF6B6B;
    transform: scale(1.05);
}
.message-container {
    background-color: #262730;
    border-radius: 10px;
    padding: 15px;
    margin-bottom: 10px;
}
.user-message {
    background-color: #4A4A4A;
    text-align: right;
}
.ai-message {
    background-color: #1E1E2E;
}
</style>
""", unsafe_allow_html=True)

# Initialize OpenAI client
apiKey = st.secrets.get("OPENAI_API_KEY")
baseURL: "https://api.x.ai/v1"

class GrokChatApp:
    def __init__(self):
        self.messages: List[Dict] = []
        self.system_prompt = (
            "You are Grok, a witty AI assistant inspired by the Hitchhiker's Guide to the Galaxy. "
            "Provide creative, humorous, and insightful responses with a touch of cosmic wisdom."
        )

    def generate_response(self, user_input: str) -> str:
        try:
            # Prepare messages for API call
            conversation = [
                {"role": "system", "content": self.system_prompt}
            ] + [
                {"role": msg['role'], "content": msg['content']} 
                for msg in self.messages
            ] + [
                {"role": "user", "content": user_input}
            ]

            # Call Grok AI
            response = openai.chat.completions.create(
                model="grok-beta",
                messages=conversation
            )

            return response.choices[0].message.content
        except Exception as e:
            return f"🚨 Error: {str(e)}"

    def add_message(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})

def main():
    st.title("🚀 Grok AI: Cosmic Conversations")
    
    # Initialize chat app
    if 'chat_app' not in st.session_state:
        st.session_state.chat_app = GrokChatApp()

    # Chat history display
    chat_container = st.container()
    with chat_container:
        for msg in st.session_state.chat_app.messages:
            with st.chat_message(msg['role'], avatar='🤖' if msg['role'] == 'assistant' else '👤'):
                st.markdown(msg['content'])

    # User input
    user_input = st.chat_input("Ask Grok anything about life, universe, and everything...")

    if user_input:
        # Display user message
        st.session_state.chat_app.add_message("user", user_input)
        with st.chat_message("user", avatar='👤'):
            st.markdown(user_input)

        # Generate and display AI response
        with st.chat_message("assistant", avatar='🤖'):
            with st.spinner("Grok is pondering the mysteries of the universe..."):
                response = st.session_state.chat_app.generate_response(user_input)
                st.markdown(response)
        
        st.session_state.chat_app.add_message("assistant", response)

if __name__ == "__main__":
    main()