Update Dockerfile

Dockerfile

@@ -1,5 +1,6 @@
 FROM ollama/ollama:latest
 
+# Install Python & Dependencies with PEP 668 fix
 RUN apt-get update && apt-get install -y python3 python3-pip && \
     pip3 install flask flask-cors requests --break-system-packages
 
@@ -8,36 +9,59 @@ ENV OLLAMA_HOST=127.0.0.1:11434
 ENV HOME=/home/ollama
 RUN mkdir -p /home/ollama/.ollama && chmod 777 /home/ollama/.ollama
 
-# Gatekeeper that manages a live whitelist
+# Create the Security Guard with Rate Limiting and Whitelisting
 RUN cat <<EOF > /guard.py
 from flask import Flask, request, Response
 import requests
 from flask_cors import CORS
+import json, os, datetime
 
 app = Flask(__name__)
 CORS(app)
 
-# Whitelist starts with one seed key
-whitelist = set(["sk-admin-seed-99"]) 
+DB_PATH = "/home/ollama/usage.json"
+LIMIT = 500
+whitelist = set(["sk-admin-seed-99"])
+
+def get_usage():
+    if not os.path.exists(DB_PATH): return {}
+    with open(DB_PATH, "r") as f: return json.load(f)
+
+def save_usage(data):
+    with open(DB_PATH, "w") as f: json.dump(data, f)
 
 @app.route("/whitelist", methods=["POST"])
 def add_to_whitelist():
     key = request.json.get("key")
     if key:
         whitelist.add(key)
-        return {"status": "Key Authorized"}, 200
+        return {"status": "Authorized"}, 200
     return {"status": "Error"}, 400
 
 @app.route("/api/generate", methods=["POST"])
 def proxy():
     user_key = request.headers.get("x-api-key", "")
     if user_key not in whitelist:
-        return {"error": "Unauthorized: This key is not registered."}, 401
+        return {"error": "Unauthorized: Key not registered"}, 401
+
+    now = datetime.datetime.now()
+    month_key = now.strftime("%Y-%m")
+    usage = get_usage()
+    
+    if user_key not in usage: usage[user_key] = {}
+    if month_key not in usage[user_key]: usage[user_key][month_key] = 0
+    
+    if usage[user_key][month_key] >= LIMIT:
+        return {"error": f"Monthly limit of {LIMIT} reached"}, 429
+    
     try:
         resp = requests.post("http://127.0.0.1:11434/api/generate", json=request.json)
+        if resp.status_code == 200:
+            usage[user_key][month_key] += 1
+            save_usage(usage)
         return Response(resp.content, resp.status_code, resp.headers.items())
     except:
-        return {"error": "Ollama starting..."}, 503
+        return {"error": "Server booting..."}, 503
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)