Update app.py

app.py

@@ -1,145 +1,202 @@
+import gradio as gr
+import requests
 import os
 import json
-import requests
-import gradio as gr
 
-FIREWORKS_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
-MODEL_ID = os.getenv("FIREWORKS_MODEL_ID", "accounts/waseem-9b447b/models/ft-gdixl08u-sz53t")
+# 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.")
 
-# Secrets (server-side only; never sent to the client UI)
-FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")  # required
-SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT")
+# API endpoint
+API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
 
-if not FIREWORKS_API_KEY:
-    raise RuntimeError("Missing FIREWORKS_API_KEY environment variable")
-
-def _fireworks_stream(payload):
-    """Generator that streams tokens from Fireworks chat completions SSE response."""
+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
+    """
+    # Build conversation history
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    
+    # Add conversation history
+    for human, assistant in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": assistant})
+    
+    # Add current message
+    messages.append({"role": "user", "content": message})
+    
+    # Prepare the request
     headers = {
         "Accept": "application/json",
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
+        "Authorization": f"Bearer {API_KEY}"
     }
-    payload = dict(payload)
-    payload["stream"] = True
-    with requests.post(FIREWORKS_URL, headers=headers, json=payload, stream=True) as r:
-        r.raise_for_status()
-        buffer = ""
-        for line in r.iter_lines(decode_unicode=True):
-            if not line:
-                continue
-            if line.startswith("data:"):
-                data = line[len("data:"):].strip()
-                if data == "[DONE]":
-                    break
-                try:
-                    obj = json.loads(data)
-                except json.JSONDecodeError:
-                    buffer += data
-                    try:
-                        obj = json.loads(buffer)
-                        buffer = ""
-                    except Exception:
-                        continue
-                try:
-                    delta = obj["choices"][0]["delta"]
-                    if "content" in delta and delta["content"]:
-                        yield delta["content"]
-                except Exception:
-                    continue
-
-def _normalize_history_to_messages(history):
-    """Normalize history from Gradio into OpenAI-style messages without system prompt."""
-    # Chatbot(type='messages') already gives a list of dicts: [{'role': 'user'|'assistant', 'content': '...'}, ...]
-    if not history:
-        return []
-    if isinstance(history, list) and len(history) > 0 and isinstance(history[0], dict) and "role" in history[0]:
-        # Already messages format; pass through (filter any roles other than user/assistant)
-        return [m for m in history if m.get("role") in ("user", "assistant")]
-    # Back-compat: history may be list of (user, assistant) tuples
-    msgs = []
-    for u, a in history:
-        if u:
-            msgs.append({"role": "user", "content": u})
-        if a:
-            msgs.append({"role": "assistant", "content": a})
-    return msgs
-
-def _build_messages(history, user_message):
-    messages = []
-    if SYSTEM_PROMPT:
-        messages.append({"role": "system", "content": SYSTEM_PROMPT})
-    messages.extend(_normalize_history_to_messages(history))
-    if user_message:
-        messages.append({"role": "user", "content": user_message})
-    return messages
-
-def chat_fn(user_message, history, max_tokens, temperature, top_p, top_k, presence_penalty, frequency_penalty):
+    
     payload = {
-        "model": MODEL_ID,
-        "max_tokens": int(max_tokens),
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "top_k": int(top_k),
-        "presence_penalty": float(presence_penalty),
-        "frequency_penalty": float(frequency_penalty),
-        "messages": _build_messages(history, user_message),
+        "model": "accounts/waseem-9b447b/models/ft-gdixl08u-sz53t",
+        "max_tokens": max_tokens,
+        "top_p": top_p,
+        "top_k": top_k,
+        "presence_penalty": 0,
+        "frequency_penalty": 0,
+        "temperature": temperature,
+        "messages": messages
     }
-    for token in _fireworks_stream(payload):
-        yield token
+    
+    try:
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
+        response.raise_for_status()
+        
+        result = response.json()
+        assistant_message = result["choices"][0]["message"]["content"]
+        return assistant_message
+        
+    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)}"
 
-def clear_history():
-    return None
+# Custom CSS for a modern look
+custom_css = """
+.gradio-container {
+    font-family: 'Inter', sans-serif;
+}
+#title {
+    text-align: center;
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    font-size: 2.5em;
+    font-weight: bold;
+    margin-bottom: 0.5em;
+}
+#description {
+    text-align: center;
+    font-size: 1.1em;
+    color: #666;
+    margin-bottom: 2em;
+}
+.message-wrap {
+    border-radius: 12px !important;
+}
+"""
 
-with gr.Blocks(theme=gr.themes.Soft(), css="""
-:root { --radius: 16px; }
-#title { font-weight: 800; letter-spacing: -0.02em; }
-div.controls { gap: 10px !important; }
-""") as demo:
-    gr.HTML("""
-    <div style="display:flex; align-items:center; gap:12px; margin: 6px 0 16px;">
-      <svg width="28" height="28" viewBox="0 0 24 24" fill="none"><path d="M12 3l7 4v6c0 5-7 8-7 8s-7-3-7-8V7l7-4z" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"/></svg>
-      <div>
-        <div id="title" style="font-size:1.25rem;">Fireworks Chat Playground</div>
-        <div style="opacity:0.7; font-size:0.95rem;">Secure, streamed chat to <code>inference/v1/chat/completions</code></div>
-      </div>
-    </div>
-    """)
+# 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.Row():
         with gr.Column(scale=3):
-            # Use messages format to avoid deprecation
-            chatbot = gr.Chatbot(height=480, type="messages", avatar_images=(None, None))
-            with gr.Row(elem_classes=["controls"]):
-                max_tokens = gr.Slider(32, 8192, value=4000, step=16, label="Max tokens")
-                temperature = gr.Slider(0.0, 2.0, value=0.6, step=0.05, label="Temperature")
-        with gr.Column(scale=2):
-            with gr.Group():
-                top_p = gr.Slider(0.0, 1.0, value=1.0, step=0.01, label="top_p")
-                top_k = gr.Slider(0, 200, value=40, step=1, label="top_k")
-                presence_penalty = gr.Slider(-2.0, 2.0, value=0.0, step=0.05, label="presence_penalty")
-                frequency_penalty = gr.Slider(-2.0, 2.0, value=0.0, step=0.05, label="frequency_penalty")
+            chatbot = gr.Chatbot(
+                height=500,
+                bubble_full_width=False,
+                avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=ai"),
+                show_copy_button=True
+            )
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    placeholder="Type your message here...",
+                    show_label=False,
+                    scale=4,
+                    container=False
+                )
+                submit_btn = gr.Button("Send 📤", scale=1, variant="primary")
+            
+            with gr.Row():
+                clear_btn = gr.ClearButton([msg, chatbot], value="Clear Chat 🗑️")
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Model Parameters")
+            
+            temperature = gr.Slider(
+                minimum=0,
+                maximum=2,
+                value=0.6,
+                step=0.1,
+                label="Temperature",
+                info="Controls randomness. Lower = more focused"
+            )
+            
+            max_tokens = gr.Slider(
+                minimum=100,
+                maximum=4000,
+                value=2000,
+                step=100,
+                label="Max Tokens",
+                info="Maximum length of response"
+            )
+            
+            top_p = gr.Slider(
+                minimum=0,
+                maximum=1,
+                value=1,
+                step=0.05,
+                label="Top P",
+                info="Nucleus sampling threshold"
+            )
+            
+            top_k = gr.Slider(
+                minimum=1,
+                maximum=100,
+                value=40,
+                step=1,
+                label="Top K",
+                info="Number of top tokens to consider"
+            )
+            
+            gr.Markdown("---")
+            gr.Markdown("### 📝 Setup Instructions")
             gr.Markdown("""
-**Security notes**
-- API key and system prompt are server-side environment variables.
-- Change the model id with `FIREWORKS_MODEL_ID` (env var).
+            1. Go to your Space **Settings**
+            2. Add these secrets:
+               - `FIREWORKS_API_KEY`: Your API key
+               - `SYSTEM_PROMPT`: Custom system prompt
+            3. Restart the Space
             """)
-            clear_btn = gr.Button("Clear", variant="secondary")
-    chat = gr.ChatInterface(
-        fn=chat_fn,
-        chatbot=chatbot,
-        additional_inputs=[max_tokens, temperature, top_p, top_k, presence_penalty, frequency_penalty],
-        title=None,
-        submit_btn="Send",
+    
+    # Handle message submission
+    def respond(message, chat_history, temp, max_tok, top_p_val, top_k_val):
+        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)
         
-examples=[
-    ["Hello!", 4000, 0.6, 1.0, 40, 0.0, 0.0],
-    ["Summarize: Why is retrieval-augmented generation useful for insurers?", 4000, 0.6, 1.0, 40, 0.0, 0.0],
-    ["Write a 3-bullet status update for the Palmyra team.", 4000, 0.6, 1.0, 40, 0.0, 0.0]
-],
-
-        description="Start chatting below. Streaming is enabled."
+        chat_history.append((message, bot_message))
+        return "", chat_history
+    
+    msg.submit(
+        respond,
+        [msg, chatbot, temperature, max_tokens, top_p, top_k],
+        [msg, chatbot]
+    )
+    
+    submit_btn.click(
+        respond,
+        [msg, chatbot, temperature, max_tokens, top_p, top_k],
+        [msg, chatbot]
     )
-    clear_btn.click(fn=clear_history, outputs=chatbot)
+    
+    # Add examples
+    gr.Examples(
+        examples=[
+            ["Hello! Can you introduce yourself?"],
+            ["What can you help me with?"],
+            ["Tell me an interesting fact about AI."],
+        ],
+        inputs=msg,
+        label="💡 Try these examples"
+    )
+    
+    gr.Markdown("""
+    ---
+    ### 🔒 Privacy & Security
+    - Your API key is stored securely in Hugging Face Spaces secrets
+    - System prompt is hidden from users
+    - All conversations are private to your session
+    """)
 
+# Launch the app
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0")
+    demo.launch()