Update app.py

app.py

@@ -1,98 +1,171 @@
 # yellowflash_with_perplexity.py
 # TEST ONLY: hardcoded keys included (do NOT publish)
 
-import time, traceback, requests
+import time
+import traceback
+import requests
 import gradio as gr
 
+# New Imports for Streaming
+from google import genai
+from groq import Groq
+from groq.types.chat import ChatCompletionChunk
+
 # ---------------------------
 # HARDCODED KEYS (TESTING)
 # ---------------------------
 GEMINI_KEY = "AIzaSyAPfDiu2V_aD6un00qHt5bkISm6C0Pkx7o"
-GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+# NOTE: Gemini Python SDK handles the URL internally for streaming
+# GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
 
 GROQ_KEY = "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"
-GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
+# NOTE: Groq Python SDK handles the URL internally
+GROQ_URL = "https://api.groq.com/openai/v1/chat/completions" 
 GROQ_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 
-# ---------------------------
-# Helpers
-# ---------------------------
-def post_with_retries(url, headers, payload, timeout=18, max_retries=2):
-    for i in range(max_retries):
-        try:
-            r = requests.post(url, headers=headers, json=payload, timeout=timeout)
-            r.raise_for_status()
-            return r
-        except Exception as e:
-            if i == max_retries - 1:
-                raise
-            time.sleep(0.5 + i)
-    raise Exception("Max retries exceeded")
+# Initialize Clients
+try:
+    # Use the hardcoded key directly in the client for testing
+    GEMINI_CLIENT = genai.Client(api_key=GEMINI_KEY)
+    GROQ_CLIENT = Groq(api_key=GROQ_KEY)
+except Exception as e:
+    # This prevents the app from crashing if keys are completely invalid/missing
+    print(f"WARNING: Failed to initialize one or both API clients: {e}")
+    GEMINI_CLIENT = None
+    GROQ_CLIENT = None
 
 # ---------------------------
-# Model callers
+# Model callers (UPDATED FOR STREAMING)
 # ---------------------------
-def call_gemini(api_key, message, history):
-    headers = {"Content-Type": "application/json", "x-goog-api-key": api_key}
+def call_gemini_stream(message, history):
+    # History format conversion for the SDK
     contents = []
-    for u, m in history:
-        contents.append({"role":"user","parts":[{"text":u}]})
-        contents.append({"role":"model","parts":[{"text":m}]})
-    contents.append({"role":"user","parts":[{"text":message}]})
-    payload = {"contents": contents}
-    r = post_with_retries(GEMINI_URL, headers, payload)
-    data = r.json()
-    return data.get("candidates",[{}])[0].get("content",{}).get("parts",[{}])[0].get("text","")
-
-def call_llama_via_groq(api_key, model, message, history):
-    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    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}]})
+
+    if GEMINI_CLIENT is None:
+        yield "Error: Gemini client not initialized. Check API key."
+        return
+
+    # Use the streaming method and yield chunks
+    response_stream = GEMINI_CLIENT.models.generate_content_stream(
+        model="gemini-2.0-flash", # Using the name directly
+        contents=contents
+    )
+    
+    full_response = ""
+    for chunk in response_stream:
+        # Check if text content exists and is not empty
+        if chunk.text:
+            full_response += chunk.text
+            yield full_response
+
+def call_llama_via_groq_stream(message, history):
+    # History format conversion to OpenAI/Groq messages
     msgs = []
-    for u, m in history:
-        msgs.append({"role":"user","content":u})
-        msgs.append({"role":"assistant","content":m})
-    msgs.append({"role":"user","content":message})
-    payload = {"model": model, "messages": msgs}
-    r = post_with_retries(GROQ_URL, headers, payload)
-    data = r.json()
-    if "choices" in data and data["choices"]:
-        ch = data["choices"][0]
-        if isinstance(ch.get("message"), dict):
-            return ch["message"].get("content","")
-        return ch.get("text","")
-    return str(data)
+    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})
+    
+    if GROQ_CLIENT is None:
+        yield "Error: Groq client not initialized. Check API key."
+        return
+
+    # Use the streaming method with stream=True
+    response_stream = GROQ_CLIENT.chat.completions.create(
+        model=GROQ_MODEL,
+        messages=msgs,
+        stream=True # Critical for streaming
+    )
+    
+    full_response = ""
+    for chunk in response_stream:
+        # Check if the chunk contains content delta
+        if isinstance(chunk, ChatCompletionChunk) and chunk.choices:
+            content = chunk.choices[0].delta.content
+            if content:
+                full_response += content
+                yield full_response
 
 # ---------------------------
-# Chat function
+# Chat function (UPDATED to be a GENERATOR)
 # ---------------------------
 def chat_fn(message, history, model_choice):
     try:
         if model_choice == "Google Gemini 2.0 Flash":
-            return call_gemini(GEMINI_KEY, message, history)
+            # Delegate to the streaming function
+            yield from call_gemini_stream(message, history) 
         elif model_choice == "Meta LLaMA 4":
-            return call_llama_via_groq(GROQ_KEY, GROQ_MODEL, message, history)
+            # Delegate to the streaming function
+            yield from call_llama_via_groq_stream(message, history)
         else:
-            return f"Unknown model: {model_choice}"
+            yield f"Unknown model: {model_choice}"
     except Exception as e:
-        return f"Error: {e}\n{traceback.format_exc()}"
+        yield f"Error: {e}\n{traceback.format_exc()}"
 
 # ---------------------------
-# Dark Mode CSS (your original)
+# CSS (UPDATED for smoother dropdown look)
 # ---------------------------
 css = """
-/* topbar layout */
+/* Topbar layout */
 #topbar { display:flex; justify-content:space-between; align-items:center;
           padding:18px 28px; background:#0f0f0f; border-bottom:1px solid #1f1f1f; }
 #title { font-weight:800; color:#ffcc33; font-size:20px; }
 
-/* compact, flat dropdown look */
-#model_dropdown .gr-dropdown { background:#1a1a1a !important; border:1px solid #2b2b2b !important;
-    color:#ddd !important; padding:10px 12px !important; border-radius:8px !important; 
-    width:260px !important; box-shadow:none !important; }
+/* ------------------------------------- */
+/* ** DROPDOWN STYLING FOR SMOOTH LOOK ** */
+/* ------------------------------------- */
+#model_dropdown .gr-dropdown-wrap {
+    /* Hide the default arrow container for a cleaner look */
+    border: none !important;
+    box-shadow: none !important;
+}
+
+#model_dropdown .gr-dropdown-wrap > button {
+    /* Main button styling (the collapsed box) */
+    background:#1a1a1a !important; 
+    border:1px solid #2b2b2b !important;
+    color:#ddd !important; 
+    padding:10px 15px !important; 
+    border-radius:10px !important;  
+    box-shadow:none !important;
+    /* Smooth transition for hover effect */
+    transition: all 0.2s ease-in-out;
+}
+#model_dropdown .gr-dropdown-wrap > button:hover {
+    border-color: #ffcc33 !important;
+    background: #252525 !important;
+}
+
+/* Style for the options list when opened */
+.gradio-container .gr-dropdown-options {
+    background: #1a1a1a !important;
+    border: 1px solid #ffcc33 !important; /* Accent border when open */
+    border-radius: 10px !important;
+    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.5);
+    padding: 5px 0;
+}
+.gradio-container .gr-dropdown-options button {
+    /* Style for individual options */
+    color: #ddd !important;
+    padding: 10px 15px !important;
+    background: transparent !important;
+    transition: background 0.1s ease;
+}
+.gradio-container .gr-dropdown-options button:hover {
+    /* Highlight the selected option with the yellow accent */
+    background: #ffcc3320 !important; /* Light yellow background on hover */
+    color: #fff !important;
+}
+/* ------------------------------------- */
 
 /* make ChatInterface chat area taller */
 .gradio-container .chat-interface .chatbot { min-height: calc(100vh - 220px); background:#111; color:#eee; }
 
-/* style send button */
+/* style send button (your original) */
 .gr-button { border-radius:10px !important; background:#2c2c3f !important; color:#fff !important; }
 """
 
@@ -105,15 +178,20 @@ with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
             choices=["Google Gemini 2.0 Flash", "Meta LLaMA 4"],
             value="Google Gemini 2.0 Flash",
             show_label=False,
-            elem_id="model_dropdown"
+            elem_id="model_dropdown",
+            scale=1 # Use scale to control width in the Row
         )
     
-
+    # We must explicitly set title and description here as per your original request 
     gr.ChatInterface(
         fn=chat_fn,
         title="⚡ YellowFlash.ai",
         description="under development",
         additional_inputs=[model_dropdown],
+        # When streaming, the Submit button automatically becomes a Stop button
     )
 
-app.launch(share=True)
+if __name__ == "__main__":
+    # Remove share=True for Hugging Face deployment
+    # The launch() call is kept simple for the HF environment
+    app.launch()