Update app.py

app.py

@@ -33,12 +33,78 @@ async def reroute_to_status():
 
 OLLAMA_LIBRARY_URL = "https://ollama.com/library"
 
-RATE_LIMIT = 60
-WINDOW_SECONDS = 60 * 60 * 24
-ip_store = {}  # { ip: { "count": int, "reset": timestamp } }
-AUDIO_RATE_LIMIT = 10
-AUDIO_WINDOW_SECONDS = 60 * 60 * 24
-audio_ip_store = {}
+PLAN_ORDER = ["free", "light", "pro", "creator", "professional"]
+TIER_CONFIG = {
+    "free": {
+        "name": "Free Tier",
+        "url": "",
+        "price": "0.00",
+        "limits": {
+            "cloudChatDaily": 50,
+            "imagesDaily": 10,
+            "videosDaily": 3,
+            "audioWeekly": 1,
+        },
+    },
+    "light": {
+        "name": "InferencePort AI Light",
+        "url": "https://buy.stripe.com/test_6oUcN5g665rp7nLgaq8bS00",
+        "price": "9.99",
+        "limits": {
+            "cloudChatDaily": None,
+            "imagesDaily": 50,
+            "videosDaily": 10,
+            "audioWeekly": 5,
+        },
+    },
+    "pro": {
+        "name": "InferencePort AI Pro",
+        "url": "https://buy.stripe.com/test_bJe9AT2fg6vt23rgaq8bS01",
+        "price": "15.99",
+        "limits": {
+            "cloudChatDaily": None,
+            "imagesDaily": 150,
+            "videosDaily": None,
+            "audioWeekly": 25,
+        },
+    },
+    "creator": {
+        "name": "InferencePort AI Creator",
+        "url": "https://buy.stripe.com/test_14AaEX9HIdXV8rPf6m8bS02",
+        "price": "29.99",
+        "limits": {
+            "cloudChatDaily": None,
+            "imagesDaily": 300,
+            "videosDaily": 50,
+            "audioWeekly": 45,
+        },
+    },
+    "professional": {
+        "name": "InferencePort AI Professional",
+        "url": "https://buy.stripe.com/test_5kQ00jf22cTR0ZncYe8bS03",
+        "price": "99.99",
+        "limits": {
+            "cloudChatDaily": None,
+            "imagesDaily": None,
+            "videosDaily": None,
+            "audioWeekly": 75,
+        },
+    },
+}
+USAGE_PERIODS = {
+    "cloudChatDaily": "daily",
+    "imagesDaily": "daily",
+    "videosDaily": "daily",
+    "audioWeekly": "weekly",
+}
+usage_store = {
+    "cloudChatDaily": {},
+    "imagesDaily": {},
+    "videosDaily": {},
+    "audioWeekly": {},
+}
+IDENTITY_CACHE_TTL_SECONDS = 60
+identity_cache = {}
 
 REASONING_KEYWORDS = [
     # explicit reasoning requests
@@ -211,28 +277,106 @@ def extract_user_text(messages: list) -> str:
         if m.get("role") == "user"
     ).lower()
 
-def check_audio_rate_limit(ip: str):
+def normalize_plan_key(plan_name: Optional[str]) -> str:
+    if not plan_name:
+        return "free"
+    normalized = "".join(ch for ch in str(plan_name).lower() if ch.isalpha())
+    if "professional" in normalized:
+        return "professional"
+    if "creator" in normalized:
+        return "creator"
+    if "pro" in normalized:
+        return "pro"
+    if "light" in normalized:
+        return "light"
+    return "free"
+
+
+def get_usage_period_key(metric: str) -> str:
+    now = time.gmtime()
+    period = USAGE_PERIODS.get(metric, "daily")
+    if period == "weekly":
+        iso_year, iso_week, _ = time.strftime("%G %V %u", now).split(" ")
+        return f"{iso_year}-W{iso_week}"
+    return time.strftime("%Y-%m-%d", now)
+
+
+async def resolve_rate_limit_identity(
+    request: Request,
+    authorization: Optional[str],
+) -> tuple[str, str]:
     now = time.time()
+    default_subject = f"ip:{request.client.host if request.client else 'unknown'}"
+    if not authorization or not authorization.startswith("Bearer "):
+        return "free", default_subject
+
+    token = authorization.split(" ", 1)[1].strip()
+    if not token:
+        return "free", default_subject
+
+    cached = identity_cache.get(token)
+    if cached and cached.get("expires_at", 0) > now:
+        return cached.get("plan_key", "free"), cached.get("subject", default_subject)
+
+    try:
+        sub = await fetch_subscription(token)
+    except Exception:
+        return "free", default_subject
+
+    if not isinstance(sub, dict) or sub.get("error"):
+        return "free", default_subject
+
+    email = sub.get("email")
+    if isinstance(email, str) and email.strip():
+        subject = f"user:{email.strip().lower()}"
+    else:
+        subject = default_subject
+
+    plan_key = normalize_plan_key(sub.get("plan_key"))
+    identity_cache[token] = {
+        "plan_key": plan_key,
+        "subject": subject,
+        "expires_at": now + IDENTITY_CACHE_TTL_SECONDS,
+    }
+    return plan_key, subject
 
-    if ip not in audio_ip_store:
-        audio_ip_store[ip] = {
-            "count": 0,
-            "reset": now + AUDIO_WINDOW_SECONDS
-        }
 
-    entry = audio_ip_store[ip]
+async def enforce_rate_limit(
+    request: Request,
+    authorization: Optional[str],
+    metric: str,
+) -> Dict[str, Optional[int | str]]:
+    if metric not in usage_store:
+        raise HTTPException(status_code=500, detail=f"Unknown limit metric: {metric}")
 
-    if now > entry["reset"]:
-        entry["count"] = 0
-        entry["reset"] = now + AUDIO_WINDOW_SECONDS
+    plan_key, subject = await resolve_rate_limit_identity(request, authorization)
+    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
+    plan_limits = plan.get("limits", {})
+    limit = plan_limits.get(metric)
 
-    if entry["count"] >= AUDIO_RATE_LIMIT:
+    if limit is None:
+        return {"plan_key": plan_key, "remaining": None}
+
+    window_key = get_usage_period_key(metric)
+    bucket = usage_store[metric]
+    entry = bucket.get(subject)
+    if not entry or entry.get("window") != window_key:
+        entry = {"window": window_key, "count": 0}
+        bucket[subject] = entry
+
+    if entry["count"] >= int(limit):
         raise HTTPException(
             status_code=429,
-            detail="Daily audio limit reached: 10 per IP"
+            detail=f"{metric} limit reached for {plan.get('name', 'current plan')}",
         )
 
     entry["count"] += 1
+    remaining = max(0, int(limit) - entry["count"])
+    return {"plan_key": plan_key, "remaining": remaining}
+
+
+async def check_audio_rate_limit(request: Request, authorization: Optional[str]):
+    await enforce_rate_limit(request, authorization, "audioWeekly")
 
 
 def is_complex_reasoning(prompt: str) -> bool:
@@ -263,35 +407,17 @@ def is_cinematic_image_prompt(prompt: str) -> bool:
             return True
     return False
 
-def check_rate_limit(ip: str):
-    now = time.time()
-
-    if ip not in ip_store:
-        ip_store[ip] = {"count": 0, "reset": now + WINDOW_SECONDS}
-
-    entry = ip_store[ip]
-
-    if now > entry["reset"]:
-        entry["count"] = 0
-        entry["reset"] = now + WINDOW_SECONDS
+async def check_image_rate_limit(request: Request, authorization: Optional[str]):
+    await enforce_rate_limit(request, authorization, "imagesDaily")
 
-    if entry["count"] >= RATE_LIMIT:
-        raise HTTPException(
-            status_code=429,
-            detail="Daily limit reached: 25 images per IP"
-        )
 
-    entry["count"] += 1
+async def check_video_rate_limit(request: Request, authorization: Optional[str]):
+    await enforce_rate_limit(request, authorization, "videosDaily")
 
 PKEY  = os.getenv("POLLINATIONS_KEY", "")
 PKEY2 = os.getenv("POLLINATIONS2_KEY", "")
 PKEY3 = os.getenv("POLLINATIONS3_KEY", "")
 
-CHAT_RATE_LIMIT = 50
-CHAT_WINDOW_SECONDS = 60 * 60
-
-chat_ip_store = {}
-
 GROQ_TOOL_MODELS = [
     "openai/gpt-oss-120b",
     "openai/gpt-oss-20b",
@@ -316,29 +442,8 @@ CEREBRAS_MODELS = [
     "zai-glm-4.7",
 ]
 
-def check_chat_rate_limit(ip: str):
-    now = time.time()
-
-    if ip not in chat_ip_store:
-        chat_ip_store[ip] = {
-            "count": 0,
-            "reset": now + CHAT_WINDOW_SECONDS
-        }
-
-    entry = chat_ip_store[ip]
-
-    if now > entry["reset"]:
-        entry["count"] = 0
-        entry["reset"] = now + CHAT_WINDOW_SECONDS
-
-    if entry["count"] >= CHAT_RATE_LIMIT:
-        raise HTTPException(
-            status_code=429,
-            detail="Chat rate limit exceeded"
-        )
-
-    entry["count"] += 1
-    return entry["count"]
+async def check_chat_rate_limit(request: Request, authorization: Optional[str]):
+    return await enforce_rate_limit(request, authorization, "cloudChatDaily")
 
 @app.head("/status/sfx")
 async def head_sfx():
@@ -423,9 +528,12 @@ async def get_status():
 
 @app.post("/gen/image")
 @app.get("/genimg/{prompt}")
-async def generate_image(request: Request, prompt: str = None):
-    client_ip = request.client.host
-    check_rate_limit(client_ip)
+async def generate_image(
+    request: Request,
+    prompt: str = None,
+    authorization: Optional[str] = Header(None),
+):
+    await check_image_rate_limit(request, authorization)
     timeout = httpx.Timeout(300.0, read=300.0)
     if prompt is None:
         prompt = (await request.json()).get("prompt")
@@ -490,14 +598,13 @@ async def get_models() -> List[Dict]:
     return models
 
 @app.post("/gen/chat/completions")
-async def generate_text(request: Request):
+async def generate_text(request: Request, authorization: Optional[str] = Header(None)):
     body = await request.json()
     messages = body.get("messages", [])
     if not isinstance(messages, list) or len(messages) == 0:
         raise HTTPException(400, "messages[] is required")
 
-    ip = request.client.host
-    msg_count = check_chat_rate_limit(ip)
+    await check_chat_rate_limit(request, authorization)
     prompt_text = extract_user_text(messages)
     
     uses_tools = (
@@ -645,9 +752,12 @@ async def generate_text(request: Request):
 
 @app.get("/gen/sfx/{prompt}")
 @app.post("/gen/sfx")
-async def gensfx(request: Request, prompt: str = None):
-    client_ip = request.client.host
-    check_audio_rate_limit(client_ip)
+async def gensfx(
+    request: Request,
+    prompt: str = None,
+    authorization: Optional[str] = Header(None),
+):
+    await check_audio_rate_limit(request, authorization)
     if prompt is None:
         prompt = (await request.json()).get("prompt")
     url = f"https://gen.pollinations.ai/audio/{prompt}?model=elevenmusic&key={PKEY}"
@@ -675,9 +785,12 @@ async def gensfx(request: Request, prompt: str = None):
 
 @app.get("/gen/tts/{prompt}")
 @app.post("/gen/tts")
-async def gensfx(request: Request, prompt: str = None):
-    client_ip = request.client.host
-    check_rate_limit(client_ip)
+async def gensfx(
+    request: Request,
+    prompt: str = None,
+    authorization: Optional[str] = Header(None),
+):
+    await check_audio_rate_limit(request, authorization)
     if prompt is None:
         prompt = (await request.json()).get("prompt")
     url = f"https://gen.pollinations.ai/audio/{prompt}?key={PKEY3}"
@@ -706,9 +819,13 @@ async def gensfx(request: Request, prompt: str = None):
 @app.get("/gen/video/{prompt}")
 @app.post("/gen/video")
 @app.head("/gen/video")
-async def genvideo(request: Request, prompt: str = None):
-    client_ip = request.client.host
-    check_rate_limit(client_ip)
+async def genvideo(
+    request: Request,
+    prompt: str = None,
+    authorization: Optional[str] = Header(None),
+):
+    if request.method != "HEAD":
+        await check_video_rate_limit(request, authorization)
     return RedirectResponse(url="/gen/video/airforce", status_code=status.HTTP_307_TEMPORARY_REDIRECT)
 
     if prompt is None:
@@ -801,7 +918,11 @@ MAX_VIDEO_RETRIES = 6
 @app.get("/gen/video/airforce/{prompt}")
 @app.post("/gen/video/airforce")
 @app.head("/gen/video/airforce")
-async def genvideo_airforce(request: Request, prompt: str = None):
+async def genvideo_airforce(
+    request: Request,
+    prompt: str = None,
+    authorization: Optional[str] = Header(None),
+):
     if request.method == "HEAD":
         return Response(
             status_code=200,
@@ -833,6 +954,8 @@ async def genvideo_airforce(request: Request, prompt: str = None):
             }
         )
 
+    await check_video_rate_limit(request, authorization)
+
     aspectRatio = "3:2"
     inputMode = "normal"
 
@@ -932,36 +1055,61 @@ async def get_subscription(authorization: Optional[str] = Header(None)):
     if "error" in result:
         raise HTTPException(401, result["error"])
 
+    plan_key = normalize_plan_key(result.get("plan_key"))
+    result["plan_key"] = plan_key
+    result["plan_name"] = (TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"])["name"]
     return result
 
-@app.get("/tiers")
-async def tiers():
+@app.get("/tier-config")
+async def tier_config():
+    plans = []
+    for idx, key in enumerate(PLAN_ORDER):
+        plan = TIER_CONFIG.get(key)
+        if not plan:
+            continue
+        plans.append(
+            {
+                "key": key,
+                "name": plan["name"],
+                "url": plan["url"],
+                "price": plan["price"],
+                "limits": plan["limits"],
+                "order": idx,
+            }
+        )
+
     return JSONResponse(
         status_code=200,
-        content=[
-            {
-                "name": "InferencePort AI Light",
-                "url": "https://buy.stripe.com/test_6oUcN5g665rp7nLgaq8bS00",
-                "price": "9.99"
-            },
-            {
-                "name": "InferencePort AI Pro",
-                "url": "https://buy.stripe.com/test_bJe9AT2fg6vt23rgaq8bS01",
-                "price": "15.99"
-            },
-            {
-                "name": "InferencePort AI Creator",
-                "url": "https://buy.stripe.com/test_14AaEX9HIdXV8rPf6m8bS02",
-                "price": "29.99"
-            },
+        content={
+            "defaultPlanKey": "free",
+            "plans": plans,
+        },
+    )
+
+@app.get("/tiers")
+async def tiers():
+    paid_plans = []
+    for key in PLAN_ORDER:
+        if key == "free":
+            continue
+        plan = TIER_CONFIG.get(key)
+        if not plan:
+            continue
+        paid_plans.append(
             {
-                "name": "InferencePort AI Professional",
-                "url": "https://buy.stripe.com/test_5kQ00jf22cTR0ZncYe8bS03",
-                "price": "99.99"
+                "key": key,
+                "name": plan["name"],
+                "url": plan["url"],
+                "price": plan["price"],
+                "limits": plan["limits"],
             }
-        ]
+        )
+
+    return JSONResponse(
+        status_code=200,
+        content=paid_plans,
     )
 
 @app.get("/portal")
 def a():
-    return RedirectResponse(url="https://billing.stripe.com/p/login/test_6oUcN5g665rp7nLgaq8bS00", status_code=status.HTTP_302_FOUND)
+    return RedirectResponse(url="https://billing.stripe.com/p/login/test_6oUcN5g665rp7nLgaq8bS00", status_code=status.HTTP_302_FOUND)