Update app.py

app.py

@@ -1,91 +1,172 @@
+# app.py
 import os
+import time
 import streamlit as st
 from groq import Groq
+from streamlit_lottie import st_lottie
+import json
+import base64
 
-# Configuration
+# --- Constants ---
 PRIMARY_MODEL = "qwen-2.5-coder-32b"
-BACKUP_MODEL = "llama3-70b-8192"
+SUPPORTED_LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
+LOTTIE_ANIMATION = "https://assets3.lottiefiles.com/packages/lf20_hi95bvmx/ai.json"
 
-# Streamlit UI Config
+# --- Setup ---
 st.set_page_config(
-    page_title="AI Code Generator",
+    page_title="CodeForge AI",
     page_icon="🚀",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-# Custom CSS for Neon Theme
-def inject_custom_css():
-    st.markdown("""
+# --- Lottie Animation Loader ---
+def load_lottie(url):
+    with open("animation.json") as f:
+        return json.load(f)
+
+# --- Custom CSS with Animations ---
+def inject_css():
+    st.markdown(f"""
     <style>
-        .main { background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); }
-        .stTextInput input { color: #4fffb0 !important; background: rgba(0,0,0,0.3) !important; }
-        .stButton>button { background: rgba(79,255,176,0.2) !important; border: 2px solid #4fffb0 !important; }
-        .code-editor { border: 2px solid #4fffb0; border-radius: 10px; padding: 1rem; background: rgba(0,0,0,0.5); }
+        @keyframes glow {{
+            0% {{ box-shadow: 0 0 5px #4fffb0; }}
+            50% {{ box-shadow: 0 0 20px #4fffb0; }}
+            100% {{ box-shadow: 0 0 5px #4fffb0; }}
+        }}
+        
+        .main {{
+            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
+            color: #ffffff;
+        }}
+        
+        .stTextInput input {{
+            animation: glow 2s infinite;
+            background: rgba(0,0,0,0.3) !important;
+            border: 2px solid #4fffb0 !important;
+        }}
+        
+        .code-editor {{
+            border: 2px solid #4fffb0;
+            border-radius: 15px;
+            padding: 1rem;
+            background: rgba(0,0,0,0.5) !important;
+            animation: glow 2s infinite;
+        }}
+        
+        .download-btn {{
+            background: linear-gradient(45deg, #4fffb0, #2bd2ff) !important;
+            color: #1a1a2e !important;
+            border: none !important;
+            border-radius: 25px;
+            padding: 12px 24px;
+            font-weight: bold;
+            transition: transform 0.3s !important;
+        }}
+        
+        .download-btn:hover {{
+            transform: scale(1.05);
+        }}
     </style>
     """, unsafe_allow_html=True)
 
-inject_custom_css()
+# --- File Download Generator ---
+def get_download_link(code, filename):
+    b64 = base64.b64encode(code.encode()).decode()
+    return f'<a href="data:file/txt;base64,{b64}" download="{filename}" class="download-btn">⬇️ Download {filename}</a>'
 
-# Initialize Groq Client
-def get_groq_client():
-    api_key = os.getenv("GROQ_API_KEY")
-    if not api_key:
-        st.error("GROQ_API_KEY not found! Please set it in environment variables or secrets.")
-        return None
-    return Groq(api_key=api_key)
+# --- Model Client ---
+@st.cache_resource
+def init_groq():
+    return Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-# Code Generation Functions
-def generate_code(query, model, client):
+# --- Code Generation Core ---
+def generate_code(client, query, lang):
     try:
-        completion = client.chat.completions.create(
-            model=model,
-            messages=[{
-                "role": "user",
-                "content": f"Generate Python code for: {query}. Include comments and error handling."
-            }],
-            temperature=0.6,
+        system_prompt = f"""
+        You are a {lang} expert. Generate:
+        1. Clean, production-ready code
+        2. With error handling
+        3. Proper documentation
+        4. Optimized for performance
+        """
+        
+        response = client.chat.completions.create(
+            model=PRIMARY_MODEL,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": query}
+            ],
+            temperature=0.7,
             max_tokens=4096,
             top_p=0.95
         )
-        return completion.choices[0].message.content
+        return response.choices[0].message.content
     except Exception as e:
-        st.error(f"Error generating code: {str(e)}")
+        st.error(f"🚨 Generation Error: {str(e)}")
         return None
 
-# Main App Interface
+# --- UI Components ---
 def main():
-    client = get_groq_client()
-    if not client:
-        return
-
-    st.title("AI Code Generator 🚀")
+    inject_css()
+    client = init_groq()
     
-    col1, col2 = st.columns([3, 1])
+    # --- Hero Section ---
+    col1, col2 = st.columns([1, 2])
     with col1:
-        query = st.text_input("Describe your coding requirement:")
-    
-    if 'generated_code' not in st.session_state:
-        st.session_state.generated_code = None
-    
-    if col2.button("🔄 New Code"):
-        st.session_state.generated_code = None
+        st_lottie(load_lottie(LOTTIE_ANIMATION), height=200)
+    with col2:
+        st.markdown("# CodeForge AI")
+        st.markdown("### _Transform Ideas into Production Code_ ⚡")
+
+    # --- Language Selector ---
+    lang = st.sidebar.selectbox("💻 Programming Language", SUPPORTED_LANGUAGES, index=0)
+    user_query = st.text_input("🎯 Describe your coding requirement:", placeholder="e.g., 'Create a REST API with JWT authentication'")
 
-    if query and not st.session_state.generated_code:
-        with st.spinner("Generating code..."):
-            code = generate_code(query, PRIMARY_MODEL, client)
-            if not code:
-                st.warning("Trying backup model...")
-                code = generate_code(query, BACKUP_MODEL, client)
+    # --- Main Generation Flow ---
+    if st.button("🚀 Generate Code", use_container_width=True):
+        if not user_query:
+            st.warning("⚠️ Please enter a coding requirement!")
+            return
+            
+        with st.spinner(f"🧠 Crafting {lang} magic..."):
+            start_time = time.time()
+            code = generate_code(client, user_query, lang)
+            gen_time = time.time() - start_time
             
             if code:
-                st.session_state.generated_code = code
-            else:
-                st.error("Failed to generate code. Please try again.")
-
-    if st.session_state.generated_code:
-        with st.expander("Generated Code", expanded=True):
-            st.markdown(f"```python\n{st.session_state.generated_code}\n```")
+                st.session_state.code = code
+                st.balloons()
+                
+                # --- Code Display ---
+                with st.expander(f"✨ Generated {lang} Code", expanded=True):
+                    st.code(code, language="python" if lang == "Python" else "javascript")
+                    
+                    # --- Download Button ---
+                    ext = {"Python": "py", "JavaScript": "js", "Java": "java", 
+                          "C++": "cpp", "Go": "go", "Rust": "rs"}[lang]
+                    st.markdown(get_download_link(code, f"generated_code.{ext}"), 
+                               unsafe_allow_html=True)
+                
+                # --- Performance Metrics ---
+                st.sidebar.markdown(f"""
+                ### 📊 Generation Metrics
+                ⏱️ **Time**: {gen_time:.2f}s  
+                📝 **Tokens**: {len(code.split())}  
+                🧠 **Model**: {PRIMARY_MODEL}
+                """)
+                
+                # --- Explanation Section ---
+                if st.button("📖 Explain Implementation"):
+                    with st.spinner("🔍 Analyzing code structure..."):
+                        explanation = client.chat.completions.create(
+                            model=PRIMARY_MODEL,
+                            messages=[{
+                                "role": "user",
+                                "content": f"Explain this {lang} code:\n\n{code}"
+                            }]
+                        )
+                        st.markdown(f"## Code Breakdown\n{explanation.choices[0].message.content}")
 
 if __name__ == "__main__":
     main()