Update app.py

app.py

@@ -1,72 +1,81 @@
 import os
 import gradio as gr
 import requests
-from huggingface_hub import InferenceClient
+import json
 
 # 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
 
-# Check paywall access and get message/link
-def check_paywall(user_id: str, feature_name: str = "premium_generation") -> tuple:
+# Check if user is connected and authorized to use the paywall
+def check_paywall(user_id: str):
     headers = {
-        "Authorization": f"Bearer {PAYWALLS_API_KEY}"
-    }
-    payload = {
-        "user_id": user_id,
-        "feature": feature_name
+        "Authorization": f"Bearer {PAYWALLS_API_KEY}",
+        "X-Paywall-User": user_id
     }
     try:
-        resp = requests.post(f"{PAYWALLS_API_URL}/access", json=payload, headers=headers, timeout=10)
+        resp = requests.get(f"{PAYWALLS_API_URL}/user/connect?user={user_id}", headers=headers, timeout=10)
+        resp.raise_for_status()
         data = resp.json()
-        access_granted = data.get("access_granted", False)
-        message = data.get("message", "")
-        return access_granted, message
+        if data.get("connected"):
+            return True, ""
+        else:
+            # Provide user the URL to connect their account or top up
+            message = f"Please connect or top up to access. Visit: {data.get('url', 'No URL provided')}"
+            return False, message
     except Exception as e:
-        print(f"[Paywall] Error: {e}")
-        return False, ""
+        print(f"[Paywall] Error checking connection: {e}")
+        return False, "Error checking paywall connection."
 
-# Respond function
-def respond(
-    message, history, system_message, max_tokens, temperature, top_p, hf_token
-):
-    user_id = "demo_user"
-    premium, paywall_msg = check_paywall(user_id)
+# 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}",
+        "Content-Type": "application/json",
+        "X-Paywall-User": user_id
+    }
     
-    # Determine tokens allowed
-    tokens_allowed = max_tokens if premium else min(120, max_tokens)
-    upgrade_prompt = ""
+    messages = [{"role": "system", "content": system_message}] + history + [{"role": "user", "content": user_message}]
+    payload = {
+        "model": HF_MODEL,
+        "messages": messages,
+        "stream": False,
+        "user": user_id,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p
+    }
     
-    # If not premium, include top-up/sign-in links
-    if not premium and paywall_msg:
-        # Expect the API message to include clickable links
-        prompt_message = paywall_msg
-    else:
-        prompt_message = ""
+    try:
+        response = requests.post(f"{PAYWALLS_API_URL}/chat/completions", headers=headers, data=json.dumps(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."
+    except Exception as e:
+        print(f"[Paywall] Chat completion error: {e}")
+        return f"Error retrieving completion: {e}"
+
+# Full respond function used by Gradio chatbot interface
+def respond(message, history, system_message, max_tokens, temperature, top_p, hf_token):
+    user_id = "demo_user"  # Replace or dynamically assign per user session
+    connected, paywall_msg = check_paywall(user_id)
     
-    # Initiate Hugging Face inference client
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+    if not connected:
+        # User not connected. Show connection URL or message.
+        return (history, paywall_msg)
     
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": message})
+    # User connected - forward chat message to paywalls chat completions endpoint
+    reply_text = chat_completion(user_id, system_message, history, message, max_tokens, temperature, top_p)
     
-    response = ""
-    for partial in client.chat_completion(
-        messages,
-        max_tokens=tokens_allowed,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = partial.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-        response += token
-        yield response + (prompt_message if not premium else "")
+    history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": reply_text}]
+    return history, ""
 
-# Gradio interface
+# Gradio interface setup
 chatbot = gr.ChatInterface(
     respond,
     type="messages",