Update app.py

app.py

@@ -1,5 +1,6 @@
-# main.py
-from fastapi import FastAPI
+import os
+import time
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 import httpx
 from bs4 import BeautifulSoup
@@ -16,6 +17,66 @@ app.add_middleware(
 
 OLLAMA_LIBRARY_URL = "https://ollama.com/library"
 
+# -----------------------------
+# RATE LIMITING (25 req/day/IP)
+# -----------------------------
+RATE_LIMIT = 25
+WINDOW_SECONDS = 60 * 60 * 24  # 24 hours
+ip_store = {}  # { ip: { "count": int, "reset": timestamp } }
+
+
+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
+
+    if entry["count"] >= RATE_LIMIT:
+        raise HTTPException(
+            status_code=429,
+            detail="Daily limit reached: 25 images per IP"
+        )
+
+    entry["count"] += 1
+
+
+# -----------------------------
+# IMAGE GENERATION ENDPOINT
+# -----------------------------
+PKEY = os.getenv("POLLINATIONS_KEY", "")  # ensure this is set in your environment
+
+
+@app.get("/genimg/{prompt}")
+async def generate_image(prompt: str, request: Request):
+    client_ip = request.client.host
+    check_rate_limit(client_ip)
+
+    url = f"https://gen.pollinations.ai/image/{prompt}?model=zimage&key={PKEY}"
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(url)
+
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Pollinations error: {response.status_code}"
+        )
+
+    return Response(
+        content=response.content,
+        media_type=response.headers.get("content-type", "image/png")
+    )
+
+
+# -----------------------------
+# EXISTING MODELS SCRAPER
+# -----------------------------
 @app.get("/models")
 async def get_models() -> List[Dict]:
     async with httpx.AsyncClient() as client:
@@ -45,6 +106,4 @@ async def get_models() -> List[Dict]:
             "link": link.get("href") if link else None,
         })
 
-    print(models)
-
     return models