Update app.py

app.py

@@ -5,14 +5,14 @@ import json
 
 # These will be set as Hugging Face Spaces secrets
 API_KEY = os.environ.get("FIREWORKS_API_KEY", "")
-SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", "You are a helpful AI assistant.")
+SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT")
 
 # API endpoint
 API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
 
 def chat_with_model(message, history, temperature, max_tokens, top_p, top_k):
     """
-    Send a message to the Fireworks AI API and return the response
+    Stream responses from the API
     """
     # Build conversation history
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
@@ -27,7 +27,7 @@ def chat_with_model(message, history, temperature, max_tokens, top_p, top_k):
     
     # Prepare the request
     headers = {
-        "Accept": "application/json",
+        "Accept": "text/event-stream",
         "Content-Type": "application/json",
         "Authorization": f"Bearer {API_KEY}"
     }
@@ -40,83 +40,115 @@ def chat_with_model(message, history, temperature, max_tokens, top_p, top_k):
         "presence_penalty": 0,
         "frequency_penalty": 0,
         "temperature": temperature,
-        "messages": messages
+        "messages": messages,
+        "stream": True
     }
     
     try:
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=120, stream=True)
         response.raise_for_status()
         
-        result = response.json()
-        assistant_message = result["choices"][0]["message"]["content"]
-        return assistant_message
+        assistant_message = ""
+        
+        for line in response.iter_lines():
+            if line:
+                line = line.decode('utf-8')
+                if line.startswith('data: '):
+                    data = line[6:]
+                    if data == '[DONE]':
+                        break
+                    try:
+                        json_data = json.loads(data)
+                        if 'choices' in json_data and len(json_data['choices']) > 0:
+                            delta = json_data['choices'][0].get('delta', {})
+                            content = delta.get('content', '')
+                            if content:
+                                assistant_message += content
+                                yield assistant_message
+                    except json.JSONDecodeError:
+                        continue
+        
+        if not assistant_message:
+            yield "❌ No response received from the model."
         
     except requests.exceptions.RequestException as e:
-        return f"❌ Error: {str(e)}\n\nPlease check your API key in Hugging Face Spaces secrets."
-    except (KeyError, IndexError) as e:
-        return f"❌ Error parsing response: {str(e)}"
+        yield f"❌ Error: {str(e)}\n\nPlease check your API key in Hugging Face Spaces secrets."
+    except Exception as e:
+        yield f"❌ Unexpected error: {str(e)}"
 
 # Custom CSS for a modern look
 custom_css = """
 .gradio-container {
     font-family: 'Inter', sans-serif;
+    max-width: 1400px !important;
 }
 #title {
     text-align: center;
-    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
     -webkit-background-clip: text;
     -webkit-text-fill-color: transparent;
-    font-size: 2.5em;
-    font-weight: bold;
-    margin-bottom: 0.5em;
+    font-size: 3em;
+    font-weight: 800;
+    margin-bottom: 0.3em;
+    letter-spacing: -0.02em;
 }
-#description {
+#subtitle {
     text-align: center;
-    font-size: 1.1em;
-    color: #666;
+    font-size: 1.2em;
+    color: #888;
     margin-bottom: 2em;
+    font-weight: 300;
 }
 .message-wrap {
-    border-radius: 12px !important;
+    border-radius: 16px !important;
+}
+footer {
+    display: none !important;
 }
 """
 
 # Create Gradio interface
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    gr.HTML("<h1 id='title'>🚀 AI Model Playground</h1>")
-    gr.HTML("<p id='description'>Powered by Fireworks AI - Fine-tuned Model</p>")
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="purple", secondary_hue="pink")) as demo:
+    gr.HTML("<h1 id='title'>Palmyra-Sec Playground</h1>")
+    gr.HTML("<p id='subtitle'>Intelligent conversations powered by advanced language models</p>")
     
     with gr.Row():
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(
-                height=500,
+                height=550,
                 bubble_full_width=False,
-                avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=ai"),
-                show_copy_button=True
+                avatar_images=(
+                    "https://api.dicebear.com/7.x/avataaars/svg?seed=user",
+                    "https://api.dicebear.com/7.x/bottts-neutral/svg?seed=ai&backgroundColor=b6e3f4"
+                ),
+                show_copy_button=True,
+                likeable=True
             )
             
             with gr.Row():
                 msg = gr.Textbox(
-                    placeholder="Type your message here...",
+                    placeholder="Ask me anything...",
                     show_label=False,
                     scale=4,
-                    container=False
+                    container=False,
+                    lines=2
                 )
-                submit_btn = gr.Button("Send 📤", scale=1, variant="primary")
+                submit_btn = gr.Button("Send 💬", scale=1, variant="primary", size="lg")
             
             with gr.Row():
-                clear_btn = gr.ClearButton([msg, chatbot], value="Clear Chat 🗑️")
+                clear_btn = gr.ClearButton([msg, chatbot], value="🗑️ Clear", size="sm")
+                retry_btn = gr.Button("🔄 Retry", size="sm", variant="secondary")
         
         with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Model Parameters")
+            gr.Markdown("### ⚙️ Model Settings")
             
             temperature = gr.Slider(
                 minimum=0,
                 maximum=2,
                 value=0.6,
                 step=0.1,
-                label="Temperature",
-                info="Controls randomness. Lower = more focused"
+                label="🌡️ Temperature",
+                info="Creativity level"
             )
             
             max_tokens = gr.Slider(
@@ -124,8 +156,8 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 maximum=4000,
                 value=2000,
                 step=100,
-                label="Max Tokens",
-                info="Maximum length of response"
+                label="📏 Max Tokens",
+                info="Response length"
             )
             
             top_p = gr.Slider(
@@ -133,8 +165,8 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 maximum=1,
                 value=1,
                 step=0.05,
-                label="Top P",
-                info="Nucleus sampling threshold"
+                label="🎯 Top P",
+                info="Diversity control"
             )
             
             top_k = gr.Slider(
@@ -142,46 +174,108 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 maximum=100,
                 value=40,
                 step=1,
-                label="Top K",
-                info="Number of top tokens to consider"
+                label="🔢 Top K",
+                info="Token selection"
             )
             
-           
+            with gr.Accordion("📚 Quick Presets", open=False):
+                gr.Markdown("""
+                **Creative** → Temp: 0.9, Top P: 0.95  
+                **Balanced** → Temp: 0.6, Top P: 1.0  
+                **Precise** → Temp: 0.3, Top P: 0.9
+                """)
+                
+                with gr.Row():
+                    creative_btn = gr.Button("🎨 Creative", size="sm")
+                    balanced_btn = gr.Button("⚖️ Balanced", size="sm")
+                    precise_btn = gr.Button("🎯 Precise", size="sm")
+            
+            gr.Markdown("---")
+            gr.Markdown("### 🔧 Setup Guide")
+            gr.Markdown("""
+            **Space Settings → Secrets:**
+            - `FIREWORKS_API_KEY`
+            - `SYSTEM_PROMPT`
+            
+            Then restart your Space ♻️
+            """)
     
-    # Handle message submission
+    # Handle message submission with streaming
     def respond(message, chat_history, temp, max_tok, top_p_val, top_k_val):
+        if not message.strip():
+            return chat_history, ""
+            
         if not API_KEY:
-            bot_message = "⚠️ Please set FIREWORKS_API_KEY in Hugging Face Spaces secrets!"
-        else:
-            bot_message = chat_with_model(message, chat_history, temp, max_tok, top_p_val, top_k_val)
+            chat_history.append((message, "⚠️ Please configure API key in Space secrets!"))
+            return chat_history, ""
         
-        chat_history.append((message, bot_message))
-        return "", chat_history
+        chat_history.append((message, ""))
+        
+        for partial_response in chat_with_model(message, chat_history[:-1], temp, max_tok, top_p_val, top_k_val):
+            chat_history[-1] = (message, partial_response)
+            yield chat_history, ""
     
     msg.submit(
         respond,
         [msg, chatbot, temperature, max_tokens, top_p, top_k],
-        [msg, chatbot]
+        [chatbot, msg]
     )
     
     submit_btn.click(
         respond,
         [msg, chatbot, temperature, max_tokens, top_p, top_k],
-        [msg, chatbot]
+        [chatbot, msg]
+    )
+    
+    # Retry last message
+    def retry_last(chat_history, temp, max_tok, top_p_val, top_k_val):
+        if not chat_history:
+            return chat_history
+        
+        last_message = chat_history[-1][0]
+        chat_history = chat_history[:-1]
+        
+        for updated_history, _ in respond(last_message, chat_history, temp, max_tok, top_p_val, top_k_val):
+            yield updated_history
+    
+    retry_btn.click(
+        retry_last,
+        [chatbot, temperature, max_tokens, top_p, top_k],
+        [chatbot]
     )
     
+    # Preset buttons
+    def set_creative():
+        return 0.9, 0.95
+    
+    def set_balanced():
+        return 0.6, 1.0
+    
+    def set_precise():
+        return 0.3, 0.9
+    
+    creative_btn.click(set_creative, None, [temperature, top_p])
+    balanced_btn.click(set_balanced, None, [temperature, top_p])
+    precise_btn.click(set_precise, None, [temperature, top_p])
+    
     # Add examples
     gr.Examples(
         examples=[
-            ["Hello! Can you introduce yourself?"],
-            ["What can you help me with?"],
-            ["Tell me an interesting fact about AI."],
+            ["Write a creative short story about time travel"],
+            ["Explain quantum computing in simple terms"],
+            ["Help me brainstorm ideas for a mobile app"],
+            ["What are the best practices for Python code?"],
         ],
         inputs=msg,
-        label="💡 Try these examples"
+        label="💡 Example Prompts"
     )
     
-   
+    gr.Markdown("""
+    ---
+    <div style='text-align: center; color: #888; font-size: 0.9em;'>
+    🔒 <b>Privacy First</b> • All credentials stored securely • Conversations are ephemeral
+    </div>
+    """)
 
 # Launch the app
 if __name__ == "__main__":