Update app.py

app.py

@@ -1,9 +1,12 @@
 import os
 import streamlit as st
 from openai import OpenAI
-from typing import List, Dict
+from typing import List, Dict, Optional
 import random
-import base64
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 # Advanced Styling and Cosmic Theme
 st.set_page_config(
@@ -13,9 +16,8 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# Cosmic Background Generator
-def get_cosmic_background():
-    # Create a dynamic, gradient cosmic background
+def get_cosmic_background() -> str:
+    """Generate a random cosmic gradient background."""
     colors = [
         "linear-gradient(135deg, #1e2ad2, #8e2de2)",
         "linear-gradient(135deg, #ff6a00, #ee0979)",
@@ -27,86 +29,82 @@ def get_cosmic_background():
 # Advanced CSS with Cosmic Design
 st.markdown(f"""
 <style>
-body {{
-    background: {get_cosmic_background()};
-    color: #ffffff;
-    font-family: 'Orbitron', sans-serif;
-}}
-
-.stApp {{
-    background: transparent;
-}}
-
-.main-container {{
-    background: rgba(0,0,0,0.7);
-    border-radius: 15px;
-    padding: 20px;
-    backdrop-filter: blur(10px);
-}}
-
-.chat-message {{
-    background: rgba(255,255,255,0.1);
-    border-radius: 10px;
-    padding: 15px;
-    margin-bottom: 10px;
-    transition: all 0.3s ease;
-}}
-
-.chat-message:hover {{
-    transform: scale(1.02);
-    box-shadow: 0 0 20px rgba(255,255,255,0.2);
-}}
-
-.stTextInput > div > div > input {{
-    background: rgba(255,255,255,0.1);
-    color: white;
-    border: 2px solid rgba(255,255,255,0.2);
-    border-radius: 10px;
-}}
-
-.stButton > button {{
-    background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
-    color: white;
-    border: none;
-    padding: 10px 20px;
-    border-radius: 20px;
-    transition: all 0.3s ease;
-}}
-
-.stButton > button:hover {{
-    transform: scale(1.1);
-    box-shadow: 0 0 20px rgba(255,255,255,0.3);
-}}
+    body {{
+        background: {get_cosmic_background()};
+        color: #ffffff;
+        font-family: 'Orbitron', sans-serif;
+    }}
+    .stApp {{
+        background: transparent;
+    }}
+    .main-container {{
+        background: rgba(0,0,0,0.7);
+        border-radius: 15px;
+        padding: 20px;
+        backdrop-filter: blur(10px);
+    }}
+    .chat-message {{
+        background: rgba(255,255,255,0.1);
+        border-radius: 10px;
+        padding: 15px;
+        margin-bottom: 10px;
+        transition: all 0.3s ease;
+    }}
+    .chat-message:hover {{
+        transform: scale(1.02);
+        box-shadow: 0 0 20px rgba(255,255,255,0.2);
+    }}
+    .stTextInput > div > div > input {{
+        background: rgba(255,255,255,0.1);
+        color: white;
+        border: 2px solid rgba(255,255,255,0.2);
+        border-radius: 10px;
+    }}
+    .stButton > button {{
+        background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
+        color: white;
+        border: none;
+        padding: 10px 20px;
+        border-radius: 20px;
+        transition: all 0.3s ease;
+    }}
+    .stButton > button:hover {{
+        transform: scale(1.1);
+        box-shadow: 0 0 20px rgba(255,255,255,0.3);
+    }}
 </style>
 """, unsafe_allow_html=True)
 
 class AdvancedGrokChatApp:
     def __init__(self):
-        # Cosmic AI Configuration
-        self.XAI_API_KEY = "X-ai = xai-1HSpHLqxC3LnInrYpwAobgEVsjchUG0PP0adniSXWGQXwq6YfvcPto9MhsS6ouQtC4a4Dh2qqXmERgQQ"
+        self.XAI_API_KEY: Optional[str] = os.getenv('XAI_API_KEY')
         
         if not self.XAI_API_KEY:
-            st.error("🛸 Cosmic Connection Lost: API Key Missing!")
+            st.error("🛸 Cosmic Connection Lost: Please set the XAI_API_KEY environment variable!")
             st.stop()
         
-        self.client = OpenAI(
-            api_key=self.XAI_API_KEY,
-            base_url="https://api.x.ai/v1"
-        )
+        try:
+            self.client = OpenAI(
+                api_key=self.XAI_API_KEY,
+                base_url="https://api.x.ai/v1"
+            )
+        except Exception as e:
+            st.error(f"🛸 Failed to initialize OpenAI client: {str(e)}")
+            st.stop()
         
         # Advanced Personality Prompts
-        self.personality_modes = {
+        self.personality_modes: Dict[str, str] = {
             "Cosmic Philosopher": "You are a wise AI that speaks like a blend of Douglas Adams and Carl Sagan.",
             "Intergalactic Comedian": "You are a witty AI that makes jokes about the universe's absurdities.",
             "Scientific Oracle": "You provide deep scientific insights with poetic eloquence."
         }
         
-        self.current_mode = "Cosmic Philosopher"
-        self.messages: List[Dict] = []
+        self.current_mode: str = "Cosmic Philosopher"
+        self.messages: List[Dict[str, str]] = []
 
     def generate_response(self, user_input: str) -> str:
+        """Generate AI response based on user input and conversation history."""
         try:
-            # Enhanced Conversation Context
             system_prompt = (
                 f"{self.personality_modes[self.current_mode]} "
                 "Respond creatively, with depth and a touch of cosmic wonder."
@@ -121,25 +119,30 @@ class AdvancedGrokChatApp:
             response = self.client.chat.completions.create(
                 model="grok-beta",
                 messages=conversation,
-                temperature=0.7,  # More creative responses
+                temperature=0.7,
                 max_tokens=300
             )
 
             return response.choices[0].message.content
         
         except Exception as e:
-            return f"🌋 Cosmic Disruption: {str(e)}"
+            st.error(f"🌋 Cosmic Disruption: {str(e)}")
+            return "I apologize, but I'm experiencing a cosmic disturbance. Please try again."
 
-    def add_message(self, role: str, content: str):
+    def add_message(self, role: str, content: str) -> None:
+        """Add a message to the conversation history."""
         self.messages.append({"role": role, "content": content})
 
-def main():
-    st.title("🚀 Grok: Cosmic Companion")
-    
-    # Initialize chat app with session state
+def initialize_session_state() -> None:
+    """Initialize or reset the session state."""
     if 'chat_app' not in st.session_state:
         st.session_state.chat_app = AdvancedGrokChatApp()
 
+def main() -> None:
+    st.title("🚀 Grok: Cosmic Companion")
+    
+    initialize_session_state()
+
     # Sidebar with Advanced Controls
     with st.sidebar:
         st.header("🌌 Cosmic Controls")
@@ -159,9 +162,13 @@ def main():
                 st.success("Conversation reset to cosmic zero!")
         
         with col2:
-            if st.button("💾 Save Conversation"):
-                # Placeholder for conversation saving logic
-                st.info("Conversation saved to cosmic archives!")
+            if st.button("💾 Save Chat"):
+                st.download_button(
+                    label="Download Chat",
+                    data=str(st.session_state.chat_app.messages),
+                    file_name="cosmic_chat.txt",
+                    mime="text/plain"
+                )
 
     # Chat Interface
     for msg in st.session_state.chat_app.messages: