Update app.py

app.py

@@ -2,13 +2,26 @@ import os
 import gradio as gr
 import requests
 import json
+from datetime import datetime
 
 # Load secrets
 PAYWALLS_API_KEY = os.environ.get("PAYWALLS_API_KEY")
 PAYWALLS_API_URL = os.environ.get("PAYWALLS_API_URL", "https://api.paywalls.ai/v1")
-HF_MODEL = "openai/gpt-oss-20b"  # Change as needed
+HF_MODEL = "openai/gpt-oss-20b"
+
+# Free tier limits
+FREE_MAX_TOKENS = 150
+PREMIUM_MAX_TOKENS = 2048
+
+# Style presets (premium feature)
+STYLE_PRESETS = {
+    "neutral": "You are a creative writing assistant.",
+    "poetic": "You are a poetic writing assistant. Write in lyrical, expressive language with vivid imagery.",
+    "dramatic": "You are a dramatic writing assistant. Create tension, conflict, and emotional depth.",
+    "comedic": "You are a comedic writing assistant. Use humor, wit, and light-hearted tone.",
+    "horror": "You are a horror writing assistant. Build suspense and create eerie, unsettling atmospheres."
+}
 
-# Check if user is connected and authorized to use the paywall
 def check_paywall(user_id: str):
     headers = {
         "Authorization": f"Bearer {PAYWALLS_API_KEY}",
@@ -18,16 +31,11 @@ def check_paywall(user_id: str):
         resp = requests.get(f"{PAYWALLS_API_URL}/user/connect?user={user_id}", headers=headers, timeout=10)
         resp.raise_for_status()
         data = resp.json()
-        if data.get("connected"):
-            return True, ""
-        else:
-            message = f"Please connect or top up to access. Visit: {data.get('url', 'No URL provided')}"
-            return False, message
+        return data.get("connected", False), data.get('url', '')
     except Exception as e:
-        print(f"[Paywall] Error checking connection: {e}")
-        return False, "Error checking paywall connection."
+        print(f"[Paywall] Error: {e}")
+        return False, ""
 
-# Perform a chat completion request via Paywalls proxy endpoint
 def chat_completion(user_id, system_message, history, user_message, max_tokens, temperature, top_p):
     headers = {
         "Authorization": f"Bearer {PAYWALLS_API_KEY}",
@@ -47,48 +55,83 @@ def chat_completion(user_id, system_message, history, user_message, max_tokens,
     }
     
     try:
-        response = requests.post(f"{PAYWALLS_API_URL}/chat/completions", headers=headers, data=json.dumps(payload), timeout=60)
+        response = requests.post(f"{PAYWALLS_API_URL}/chat/completions", headers=headers, json=payload, timeout=60)
         response.raise_for_status()
         data = response.json()
         choices = data.get("choices", [])
         if choices:
             return choices[0]["message"]["content"]
-        else:
-            return "No response from model."
+        return "No response from model."
     except Exception as e:
-        print(f"[Paywall] Chat completion error: {e}")
-        return f"Error retrieving completion: {e}"
+        print(f"[Paywall] Error: {e}")
+        return f"Error: {e}"
 
-# Respond function updated to always return two lists (chat histories)
-def respond(
-    message, history, system_message, max_tokens, temperature, top_p, hf_token
-):
-    user_id = "demo_user"
-    connected, paywall_msg = check_paywall(user_id)
+def respond(message, history, style, max_tokens, temperature, top_p):
+    user_id = "demo_user"  # Replace with gr.Request for real user sessions
+    is_premium, paywall_url = check_paywall(user_id)
     
-    if not connected:
-        return paywall_msg  # ✅ Return just the message
+    # Apply free tier limits
+    if not is_premium:
+        max_tokens = min(max_tokens, FREE_MAX_TOKENS)
+        if style != "neutral":
+            return f"🔒 Style customization requires premium. [Unlock here]({paywall_url})"
     
-    reply_text = chat_completion(user_id, system_message, history, message, max_tokens, temperature, top_p)
-    return reply_text  # ✅ Return just the reply
+    system_message = STYLE_PRESETS.get(style, STYLE_PRESETS["neutral"])
+    reply = chat_completion(user_id, system_message, history, message, max_tokens, temperature, top_p)
+    
+    # Add tier indicator
+    tier_badge = "✨ PREMIUM" if is_premium else "🆓 FREE"
+    return f"{reply}\n\n---\n_{tier_badge}_"
 
-# Gradio chat interface with additional inputs
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a creative writing assistant.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-        gr.Textbox(value=os.environ.get("HF_TOKEN", ""), label="Hugging Face Token", visible=True, type="password"),
-    ],
-)
+def export_conversation(history):
+    user_id = "demo_user"
+    is_premium, paywall_url = check_paywall(user_id)
+    
+    if not is_premium:
+        return None, f"🔒 Export requires premium. [Unlock here]({paywall_url})"
+    
+    # Generate text file
+    content = f"Creative Writing Session - {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n"
+    for msg in history:
+        role = msg["role"].upper()
+        content += f"{role}: {msg['content']}\n\n"
+    
+    filepath = "/tmp/creative_writing.txt"
+    with open(filepath, "w") as f:
+        f.write(content)
+    
+    return filepath, "✅ Exported successfully!"
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+# Build UI
+with gr.Blocks(title="Creative Writing Assistant") as demo:
+    gr.Markdown("# 📝 Creative Writing Assistant\n**Free:** 150 tokens | **Premium:** 2048 tokens + styles + export")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(type="messages", height=500)
+            msg = gr.Textbox(label="Your prompt", placeholder="Write a story about...", lines=3)
+            submit = gr.Button("✨ Generate", variant="primary")
+            
+        with gr.Column(scale=1):
+            style = gr.Dropdown(
+                choices=list(STYLE_PRESETS.keys()),
+                value="neutral",
+                label="🎨 Writing Style (Premium)",
+                info="Unlock premium for custom styles"
+            )
+            max_tokens = gr.Slider(FREE_MAX_TOKENS, PREMIUM_MAX_TOKENS, value=FREE_MAX_TOKENS, step=50, label="Max Tokens")
+            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Creativity")
+            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Diversity")
+            
+            gr.Markdown("---")
+            export_btn = gr.Button("💾 Export (Premium)", variant="secondary")
+            download_file = gr.File(label="Download")
+            export_status = gr.Textbox(label="Status", interactive=False)
+    
+    # Event handlers
+    submit.click(respond, [msg, chatbot, style, max_tokens, temperature, top_p], [chatbot])
+    msg.submit(respond, [msg, chatbot, style, max_tokens, temperature, top_p], [chatbot])
+    export_btn.click(export_conversation, [chatbot], [download_file, export_status])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()