Update app.py

app.py

@@ -6,29 +6,26 @@ from dotenv import load_dotenv
 import os
 from groq import Groq
 
-# Load env variables
+# Load API keys
 load_dotenv()
 HF_API_KEY = st.secrets.get("HUGGINGFACE_API_KEY") or os.getenv("HUGGINGFACE_API_KEY")
 GROQ_API_KEY = st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")
 
-# Initialize Groq client if key exists
+# Initialize Groq client
 groq_client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None
 
-# Available models
-MODELS = {
-    "Python/JS Code Gen": {
-        "hf": "bigcode/starcoder2-3b",
-        "groq": "llama3-8b-8192"
-    },
-    "Debugging": {
-        "hf": "bigcode/starcoder2-3b",
-        "groq": "llama3-8b-8192"
-    },
-    "Explanation": {
-        "hf": "mistralai/Mistral-7B-Instruct-v0.2",
-        "groq": "llama3-70b-8192"
-    }
-}
+# Model lists
+HF_MODELS = [
+    "bigcode/starcoder2-3b",
+    "mistralai/Mistral-7B-Instruct-v0.2",
+    "tiiuae/falcon-7b-instruct"
+]
+
+GROQ_MODELS = [
+    "llama3-8b-8192",
+    "llama3-70b-8192",
+    "mixtral-8x7b-32768"
+]
 
 # Hugging Face API
 def query_hf(model_id, prompt, max_new_tokens=300):
@@ -71,44 +68,50 @@ def query_groq(model_id, prompt, max_tokens=300):
 
 # Prompt builder
 def build_prompt(mode, user_input, language="Python"):
-    if mode == "Python/JS Code Gen":
-        return f"You are a helpful programmer. Write {language} code for:\n{user_input}\nReturn only code."
-    elif mode == "Debugging":
+    if mode == "Generate":
+        return f"Write {language} code for this task:\n{user_input}\nOnly return code."
+    elif mode == "Debug":
         return f"Debug this code and explain briefly:\n{user_input}"
-    elif mode == "Explanation":
+    elif mode == "Explain":
         return f"Explain this code step by step:\n{user_input}"
     return user_input
 
-# Query with fallback
-def run_query(backend, mode, user_input, lang="Python"):
-    model_id = MODELS[mode][backend]
+# Run query with fallback
+def run_query(backend, model_id, mode, user_input, lang="Python"):
     prompt = build_prompt(mode, user_input, lang)
 
     if backend == "hf":
         output, latency, success = query_hf(model_id, prompt)
         if success:
-            return output, latency, "Hugging Face ✅"
+            return output, latency, f"Hugging Face ({model_id}) ✅"
         else:
             # fallback to Groq
-            output2, latency2, success2 = query_groq(MODELS[mode]["groq"], prompt)
-            return output2, latency + latency2, "Hugging Face ❌ → Groq ✅"
+            fallback_model = GROQ_MODELS[0]  # default to Llama3-8B
+            output2, latency2, success2 = query_groq(fallback_model, prompt)
+            return output2, latency + latency2, f"Hugging Face ❌ → Groq ({fallback_model}) ✅"
     elif backend == "groq":
         output, latency, success = query_groq(model_id, prompt)
         if success:
-            return output, latency, "Groq ✅"
+            return output, latency, f"Groq ({model_id}) ✅"
         else:
-            # fallback to HF
-            output2, latency2, success2 = query_hf(MODELS[mode]["hf"], prompt)
-            return output2, latency + latency2, "Groq ❌ → Hugging Face ✅"
+            # fallback to Hugging Face
+            fallback_model = HF_MODELS[0]  # default to Starcoder2
+            output2, latency2, success2 = query_hf(fallback_model, prompt)
+            return output2, latency + latency2, f"Groq ❌ → Hugging Face ({fallback_model}) ✅"
     return "⚠️ Invalid backend", 0, "None"
 
 # Streamlit UI
 st.set_page_config(page_title="CodeCraft AI", layout="wide")
-st.title("🧑‍💻 CodeCraft AI (with Fallback)")
-st.write("Generate, debug, and explain code using Hugging Face & Groq. If one fails, it auto-falls back!")
+st.title("🧑‍💻 CodeCraft AI (Model Selection + Fallback)")
+st.write("Choose models yourself. If one fails, fallback ensures smooth output.")
 
 backend = st.radio("Choose Backend", ["hf", "groq"], format_func=lambda x: "Hugging Face" if x == "hf" else "Groq")
 
+if backend == "hf":
+    model_id = st.selectbox("Choose Hugging Face Model", HF_MODELS)
+else:
+    model_id = st.selectbox("Choose Groq Model", GROQ_MODELS)
+
 tab1, tab2, tab3, tab4 = st.tabs(["Generate", "Debug", "Explain", "Analytics"])
 
 # Track logs
@@ -119,10 +122,10 @@ if "logs" not in st.session_state:
 with tab1:
     st.subheader("Code Generation")
     lang = st.selectbox("Choose language", ["Python", "JavaScript"])
-    problem = st.text_area("Enter your problem statement")
+    problem = st.text_area("Enter problem statement")
     if st.button("Generate Code", key="gen_btn"):
         if problem.strip():
-            output, latency, status = run_query(backend, "Python/JS Code Gen", problem, lang)
+            output, latency, status = run_query(backend, model_id, "Generate", problem, lang)
             st.code(output, language=lang.lower())
             st.success(f"{status} | Time: {latency:.2f}s")
             st.session_state.logs.append(("Generate", latency, status))
@@ -135,7 +138,7 @@ with tab2:
     buggy_code = st.text_area("Paste buggy code here")
     if st.button("Debug Code", key="debug_btn"):
         if buggy_code.strip():
-            output, latency, status = run_query(backend, "Debugging", buggy_code)
+            output, latency, status = run_query(backend, model_id, "Debug", buggy_code)
             st.text_area("AI Fix & Explanation", output, height=300)
             st.success(f"{status} | Time: {latency:.2f}s")
             st.session_state.logs.append(("Debug", latency, status))
@@ -148,7 +151,7 @@ with tab3:
     code_input = st.text_area("Paste code to explain")
     if st.button("Explain Code", key="explain_btn"):
         if code_input.strip():
-            output, latency, status = run_query(backend, "Explanation", code_input)
+            output, latency, status = run_query(backend, model_id, "Explain", code_input)
             st.text_area("AI Explanation", output, height=300)
             st.success(f"{status} | Time: {latency:.2f}s")
             st.session_state.logs.append(("Explain", latency, status))
@@ -162,5 +165,6 @@ with tab4:
         df = pd.DataFrame(st.session_state.logs, columns=["Mode", "Latency", "Status"])
         st.write(df)
         st.bar_chart(df.groupby("Mode")["Latency"].mean())
+        st.download_button("Download Logs as CSV", df.to_csv(index=False), "logs.csv", "text/csv")
     else:
         st.info("No usage yet. Try generating, debugging, or explaining first!")