Update app.py

app.py

@@ -1,42 +1,72 @@
-# yellowflash_with_perplexity.py
-# TEST ONLY: hardcoded keys included (do NOT publish)
+# yellowflash_with_streamlit.py
+# Production-ready structure using Streamlit
 
+import streamlit as st
 import time
 import traceback
-import gradio as gr
 
-# New Imports for Streaming (must be in your requirements.txt: google-genai, groq)
+# Import API Libraries
 from google import genai
 from groq import Groq
-from groq.types.chat import ChatCompletionChunk
-
 
 # ---------------------------
-# HARDCODED KEYS (TESTING)
+# API Key Setup (Hardcoded for current TEST request)
+# NOTE: In production, load these from st.secrets instead.
 # ---------------------------
 GEMINI_KEY = "AIzaSyAPfDiu2V_aD6un00qHt5bkISm6C0Pkx7o"
 GROQ_KEY = "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"
 GROQ_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 
-# Initialize Clients
 try:
     GEMINI_CLIENT = genai.Client(api_key=GEMINI_KEY)
     GROQ_CLIENT = Groq(api_key=GROQ_KEY)
 except Exception as e:
-    print(f"WARNING: Failed to initialize one or both API clients: {e}")
+    st.error(f"Failed to initialize API clients. Check keys. Error: {e}")
     GEMINI_CLIENT = None
     GROQ_CLIENT = None
 
+# ---------------------------
+# Streamlit Session State Initialization
+# ---------------------------
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# ---------------------------
+# Streamlit UI Setup
+# ---------------------------
+
+# Use a two-column layout for the header
+col_title, col_dropdown = st.columns([3, 1])
+
+with col_title:
+    # YellowFlash.ai Title
+    st.markdown(
+        f'<h1 style="color: #ffcc33; font-weight: 800; font-size: 24px; margin: 0;">⚡ YellowFlash.ai</h1>',
+        unsafe_allow_html=True
+    )
+
+with col_dropdown:
+    # Model Dropdown (Universal and automatically responsive)
+    model_choice = st.selectbox(
+        label="Model",
+        options=["Google Gemini 2.0 Flash", "Meta LLaMA 4"],
+        index=0,
+        label_visibility="hidden",
+        key="model_selector"
+    )
+
+st.markdown("_under development_")
+st.divider() # Clean separation line
 
 # ---------------------------
-# Model callers (STREAMING GENERATORS)
+# Model Callers (Modified for Streamlit)
 # ---------------------------
-def call_gemini_stream(message, history):
+
+def call_gemini_stream(prompt, history):
     contents = []
-    for user_msg, model_msg in history:
-        contents.append({"role": "user", "parts": [{"text": user_msg}]})
-        contents.append({"role": "model", "parts": [{"text": model_msg}]})
-    contents.append({"role": "user", "parts": [{"text": message}]})
+    # Streamlit history format is dicts: {"role": "user/assistant", "content": "..."}
+    for message in history:
+        contents.append({"role": message["role"].replace("assistant", "model"), "parts": [{"text": message["content"]}]})
 
     if GEMINI_CLIENT is None:
         yield "Error: Gemini client not initialized. Check API key."
@@ -47,18 +77,14 @@ def call_gemini_stream(message, history):
         contents=contents
     )
     
-    full_response = ""
     for chunk in response_stream:
         if chunk.text:
-            full_response += chunk.text
-            yield full_response
+            yield chunk.text
 
-def call_llama_via_groq_stream(message, history):
+def call_llama_via_groq_stream(prompt, history):
     msgs = []
-    for user_msg, model_msg in history:
-        msgs.append({"role": "user", "content": user_msg})
-        msgs.append({"role": "assistant", "content": model_msg})
-    msgs.append({"role": "user", "content": message})
+    for message in history:
+        msgs.append({"role": message["role"], "content": message["content"]})
     
     if GROQ_CLIENT is None:
         yield "Error: Groq client not initialized. Check API key."
@@ -70,155 +96,45 @@ def call_llama_via_groq_stream(message, history):
         stream=True
     )
     
-    full_response = ""
     for chunk in response_stream:
-        if isinstance(chunk, ChatCompletionChunk) and chunk.choices:
-            content = chunk.choices[0].delta.content
-            if content:
-                full_response += content
-                yield full_response
-
-# ---------------------------
-# Chat function (Main generator)
-# ---------------------------
-def chat_fn(message, history, model_choice):
-    try:
-        if model_choice == "Google Gemini 2.0 Flash":
-            yield from call_gemini_stream(message, history) 
-        elif model_choice == "Meta LLaMA 4":
-            yield from call_llama_via_groq_stream(message, history)
-        else:
-            yield f"Unknown model: {model_choice}"
-    except Exception as e:
-        yield f"Error: An API error occurred: {type(e).__name__}. Check server logs."
+        # Groq streaming provides the delta content
+        content = chunk.choices[0].delta.content
+        if content:
+            yield content
 
 # ---------------------------
-# CSS (Final Responsive Style)
+# Main Chat Logic
 # ---------------------------
-css = """
-/* The topbar is a flexible row */
-#topbar { 
-    display:flex; 
-    justify-content:flex-start; 
-    align-items:center;
-    gap: 20px; 
-    padding:18px 28px; 
-    background:#0f0f0f; 
-    border-bottom:1px solid #1f1f1f; 
-    min-height: 80px; 
-    flex-wrap: wrap; /* CRITICAL FOR MOBILE: Allows items to wrap */
-}
-
-/* Main App Title */
-#title { 
-    font-weight:800; 
-    color:#ffcc33; 
-    font-size:20px; 
-    margin: 0;
-    flex-shrink: 0; /* Prevents title from shrinking */
-}
-
-/* ------------------------------------- */
-/* ** DROPDOWN STYLING (Wider & Polished) ** */
-/* ------------------------------------- */
-
-/* Set the default width for the dropdown component */
-#model_dropdown {
-    width: 300px; 
-    min-width: 250px;
-}
-
-/* Hide Gradio's default label */
-#model_dropdown .gr-dropdown-label {
-    display: none !important;
-}
-
-/* Style the actual dropdown button */
-#model_dropdown .gr-dropdown-wrap > button {
-    background:#1a1a1a !important; 
-    border:1px solid #333 !important;
-    color:#ddd !important; 
-    padding:10px 15px !important; 
-    border-radius:10px !important;  
-    box-shadow:none !important;
-    transition: all 0.2s ease-in-out;
-    font-size: 16px; 
-    height: auto !important;
-    width: 100%; 
-}
-#model_dropdown .gr-dropdown-wrap > button:hover {
-    border-color: #ffcc33 !important;
-    background: #252525 !important;
-}
-
-/* Style the options list when opened */
-.gradio-container .gr-dropdown-options {
-    background: #1a1a1a !important;
-    border: 1px solid #ffcc33 !important;
-    border-radius: 10px !important;
-    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.5);
-    padding: 5px 0;
-}
-.gradio-container .gr-dropdown-options button:hover {
-    background: #ffcc3320 !important; 
-    color: #fff !important;
-}
-/* ------------------------------------- */
-
-
-/* Make ChatInterface/Chatbot screen adjusting */
-.gradio-container .chat-interface .chatbot { 
-    min-height: calc(100vh - 220px); 
-    background:#111; 
-    color:#eee; 
-    flex-grow: 1; 
-}
-
-/* style send button */
-.gr-button { border-radius:10px !important; background:#2c2c3f !important; color:#fff !important; }
-
-/* Mobile View Adjustments: Stack title and dropdown on small screens */
-@media (max-width: 600px) {
-    #topbar {
-        flex-direction: column; 
-        align-items: flex-start; 
-        gap: 10px;
-    }
-    #model_dropdown {
-        width: 100%; 
-        max-width: none;
-    }
-}
-"""
 
-# ---------------------------
-# Build UI (Final Layout)
-# ---------------------------
-with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
-    with gr.Row(elem_id="topbar"):
-        
-        # 1. Title (FIXED: scale argument removed)
-        # This will take up the remaining space because the dropdown is scaled to 0
-        gr.Markdown(f'<span id="title">⚡ YellowFlash.ai</span>', elem_id="title_md") 
-        
-        # 2. Dropdown 
-        model_dropdown = gr.Dropdown(
-            choices=["Google Gemini 2.0 Flash", "Meta LLaMA 4"],
-            value="Google Gemini 2.0 Flash",
-            show_label=False,
-            elem_id="model_dropdown",
-            scale=0 # Correct way to prevent it from growing horizontally
-        )
-
-    # Keep the original "under development" text
-    gr.Markdown("under development")
-
-    gr.ChatInterface(
-        fn=chat_fn,
-        title="", 
-        description="",
-        additional_inputs=[model_dropdown],
-    )
+# 1. Display chat history
+for message in st.session_state.messages:
+    # Streamlit uses st.chat_message for history display
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-if __name__ == "__main__":
-    app.launch()
+# 2. Handle user input
+if prompt := st.chat_input("Ask YellowFlash..."):
+    
+    # Add user message to state and display
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Determine which model to call
+    if model_choice == "Google Gemini 2.0 Flash":
+        stream_func = call_gemini_stream
+    else:
+        stream_func = call_llama_via_groq_stream
+
+    # Display and stream AI response
+    with st.chat_message("assistant"):
+        try:
+            # st.write_stream handles the generator function and types out the response
+            response_generator = stream_func(prompt, st.session_state.messages)
+            full_response = st.write_stream(response_generator)
+            
+            # Save the full response to state
+            st.session_state.messages.append({"role": "assistant", "content": full_response})
+            
+        except Exception as e:
+            st.error(f"An error occurred while streaming from the API: {e}")