Update app.py

app.py

@@ -1,124 +1,205 @@
 # app.py
 import os
 import time
+import json
+import zipfile
+import io
 import streamlit as st
 from groq import Groq
 from streamlit_lottie import st_lottie
-import requests
-import json
-import base64
+from streamlit_extras.stylable_container import stylable_container
 
 # --- Constants ---
 PRIMARY_MODEL = "qwen-2.5-coder-32b"
-SUPPORTED_LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
-LOTTIE_URL = "https://assets3.lottiefiles.com/packages/lf20_hi95bvmx/ai.json"
+BACKUP_MODEL = "llama3-70b-8192"
+LANGUAGES = ["Python", "JavaScript", "Java", "C++", "Go", "Rust"]
+THEMES = ["Dracula", "Monokai", "Nord", "Solarized", "One Dark"]
+
+# --- Load Assets ---
+def load_lottie(filepath):
+    with open(filepath, "r") as f:
+        return json.load(f)
+
+# --- Initialize Clients ---
+@st.cache_resource
+def init_groq():
+    return Groq(api_key=os.getenv("GROQ_API_KEY"))
 
-# --- Setup ---
+# --- UI Config ---
 st.set_page_config(
     page_title="CodeForge AI",
-    page_icon="🚀",
+    page_icon="⚡",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-# --- Lottie Animation Loader ---
-def load_lottie(url):
-    response = requests.get(url)
-    return response.json() if response.status_code == 200 else None
-
 # --- Custom CSS ---
 def inject_css():
     st.markdown(f"""
     <style>
-        .main {{
-            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
-            color: #ffffff;
+        [data-testid="stAppViewContainer"] {{
+            background: linear-gradient(135deg, #0f0c29 0%, #302b63 50%, #24243e 100%);
         }}
-        
         .stTextInput input {{
+            color: #00ffaa !important;
             background: rgba(0,0,0,0.3) !important;
-            border: 2px solid #4fffb0 !important;
-            color: #4fffb0 !important;
+            border: 2px solid #00ffaa !important;
+            box-shadow: 0 0 10px rgba(0,255,170,0.5);
         }}
-        
-        .code-editor {{
-            border: 2px solid #4fffb0;
-            border-radius: 15px;
-            padding: 1rem;
-            background: rgba(0,0,0,0.5) !important;
-        }}
-        
-        .download-btn {{
-            background: linear-gradient(45deg, #4fffb0, #2bd2ff) !important;
-            color: #1a1a2e !important;
+        .stButton>button {{
+            background: linear-gradient(90deg, #00ffaa 0%, #0088ff 100%) !important;
+            color: #111 !important;
+            font-weight: bold !important;
             border: none !important;
-            border-radius: 25px;
-            padding: 12px 24px;
+            box-shadow: 0 0 15px rgba(0,255,170,0.7);
+            transition: all 0.3s ease !important;
+        }}
+        .stButton>button:hover {{
+            transform: scale(1.05);
+            box-shadow: 0 0 20px rgba(0,255,170,0.9);
+        }}
+        .code-header {{
+            color: #00ffaa;
+            font-family: 'Fira Code', monospace;
+            text-shadow: 0 0 10px rgba(0,255,170,0.7);
+        }}
+        [data-testid="stSidebar"] {{
+            background: rgba(16, 18, 33, 0.8) !important;
+            backdrop-filter: blur(10px);
+            border-right: 1px solid rgba(0,255,170,0.3);
         }}
     </style>
     """, unsafe_allow_html=True)
 
-# --- 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>'
+inject_css()
 
-# --- Model Client ---
-@st.cache_resource
-def init_groq():
-    return Groq(api_key=os.environ.get("GROQ_API_KEY"))
+# --- Animation ---
+def render_hero():
+    col1, col2 = st.columns([2, 1])
+    with col1:
+        st.markdown("""
+        # <span class="code-header">⚡ CodeForge AI</span>
+        ### _Generate • Optimize • Deploy_
+        """, unsafe_allow_html=True)
+    with col2:
+        st_lottie(load_lottie("https://assets9.lottiefiles.com/packages/lf20_2znxgjyt.json"), 
+                 height=150, key="hero")
+
+# --- Main Generator ---
+def generate_code(client, prompt, language, model):
+    messages = [{
+        "role": "system",
+        "content": f"You are an expert {language} developer. Generate clean, efficient, production-ready code."
+    }, {
+        "role": "user",
+        "content": prompt
+    }]
+    
+    response = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        temperature=0.7,
+        max_tokens=4000,
+        top_p=0.95,
+        stream=True
+    )
+    
+    return response
+
+# --- UI Components ---
+def language_selector():
+    with st.sidebar:
+        with stylable_container(
+            key="lang-select",
+            css_styles="""
+                {
+                    border: 2px solid #00ffaa;
+                    border-radius: 10px;
+                    padding: 10px;
+                    background: rgba(0,0,0,0.2);
+                }
+            """
+        ):
+            return st.selectbox("🗝️ Select Language", LANGUAGES, index=0)
+
+def theme_selector():
+    with st.sidebar:
+        with stylable_container(
+            key="theme-select",
+            css_styles="""
+                {
+                    border: 2px solid #0088ff;
+                    border-radius: 10px;
+                    padding: 10px;
+                    background: rgba(0,0,0,0.2);
+                }
+            """
+        ):
+            return st.selectbox("🎨 Editor Theme", THEMES, index=0)
 
-# --- Code Generation Core ---
-def generate_code(client, query, lang):
-    try:
-        response = client.chat.completions.create(
-            model=PRIMARY_MODEL,
-            messages=[{
-                "role": "user",
-                "content": f"Generate {lang} code for: {query}. Include comments and error handling."
-            }],
-            temperature=0.7,
-            max_tokens=4096
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
-        return None
+# --- Download Handler ---
+def create_download_zip(code, language):
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, "a") as zf:
+        zf.writestr(f"generated_code.{language.lower()}", code)
+    return zip_buffer.getvalue()
 
-# --- Main UI ---
+# --- Main App ---
 def main():
-    inject_css()
     client = init_groq()
+    render_hero()
     
-    # --- Hero Section ---
-    col1, col2 = st.columns([1, 2])
-    with col1:
-        lottie_anim = load_lottie(LOTTIE_URL)
-        if lottie_anim:
-            st_lottie(lottie_anim, height=150)
-    with col2:
-        st.markdown("# CodeForge AI")
-        st.markdown("### _Transform Ideas into Production Code_ ⚡")
-
-    # --- Language Selector ---
-    lang = st.sidebar.selectbox("💻 Language", SUPPORTED_LANGUAGES)
-    user_query = st.text_input("🎯 Describe your requirement:", placeholder="e.g., 'Create a REST API with JWT authentication'")
-
-    # --- Generation Flow ---
-    if st.button("🚀 Generate Code", use_container_width=True):
-        if user_query:
-            with st.spinner(f"🧠 Crafting {lang} code..."):
-                code = generate_code(client, user_query, lang)
-                if code:
-                    st.session_state.code = code
-                    st.balloons()
-                    
-                    with st.expander(f"✨ Generated Code", expanded=True):
-                        st.code(code)
-                        ext = {"Python": "py", "JavaScript": "js"}.get(lang, "txt")
-                        st.markdown(get_download_link(code, f"code.{ext}"), unsafe_allow_html=True)
-        else:
-            st.warning("Please enter a requirement!")
+    # Sidebar Controls
+    language = language_selector()
+    theme = theme_selector()
+    
+    # Main Input
+    with st.form(key="input_form"):
+        prompt = st.text_area("💡 Describe your coding task:", height=150,
+                            placeholder="e.g. 'Create a REST API with FastAPI and MongoDB'")
+        
+        col1, col2, col3 = st.columns([1,1,2])
+        with col1:
+            generate_btn = st.form_submit_button("✨ Generate Code")
+        with col2:
+            explain_btn = st.form_submit_button("📚 Explain")
+        with col3:
+            optimize_btn = st.form_submit_button("⚡ Optimize")
+    
+    # Generation Logic
+    if generate_btn and prompt:
+        with st.spinner("🚀 Generating your code..."):
+            try:
+                response = generate_code(client, prompt, language, PRIMARY_MODEL)
+                
+                code_container = st.empty()
+                full_response = []
+                
+                for chunk in response:
+                    if chunk.choices[0].delta.content:
+                        full_response.append(chunk.choices[0].delta.content)
+                        current_code = "".join(full_response)
+                        code_container.code(current_code, language=language.lower(), line_numbers=True)
+                
+                st.session_state.generated_code = "".join(full_response)
+                
+                # Download Button
+                zip_data = create_download_zip(st.session_state.generated_code, language)
+                st.download_button(
+                    label="📥 Download Code",
+                    data=zip_data,
+                    file_name=f"codeforge_{int(time.time())}.zip",
+                    mime="application/zip"
+                )
+                
+            except Exception as e:
+                st.error(f"🚨 Generation failed: {str(e)}")
+                st.info("Attempting with backup model...")
+                try:
+                    response = generate_code(client, prompt, language, BACKUP_MODEL)
+                    # ... [backup handling logic]
+                except Exception as e:
+                    st.error("All models failed. Please try again later.")
 
 if __name__ == "__main__":
     main()