Adding logic to rate limit, get JWT token with user identity

app.py

@@ -9,23 +9,31 @@ import jwt
 from typing import Dict, Any
 import autopep8
 import textwrap
-
 from datasets import load_dataset
-import random
-import asyncio
+import time
+from collections import defaultdict
+import threading
+import hashlib
+
+# Rate limiting data structures
+ip_usage = defaultdict(int)
+session_usage = defaultdict(int)
+last_reset_time = time.time()
+rate_limit_lock = threading.Lock()
+
+# Constants
+MAX_IP_USAGE = 10
+MAX_SESSION_USAGE = 2
+RESET_INTERVAL = 24 * 60 * 60  # 24 hours in seconds
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Load the dataset (you might want to do this once at the start of your script)
+# Load the dataset
 dataset = load_dataset("sugiv/leetmonkey_python_dataset")
 train_dataset = dataset["train"]
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 # JWT settings
 JWT_SECRET = os.environ.get("JWT_SECRET")
 if not JWT_SECRET:
@@ -41,7 +49,9 @@ model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME, cache_dir="./
 llm = Llama(model_path=model_path, n_ctx=1024, n_threads=8, n_gpu_layers=-1, verbose=False, mlock=True)
 logger.info("8-bit model loaded successfully")
 
+# User data storage
 user_data = {}
+token_to_problem_solution = {}
 
 # Generation parameters
 generation_kwargs = {
@@ -54,6 +64,18 @@ generation_kwargs = {
     "repeat_penalty": 1.1
 }
 
+def generate_user_identifier(request: gr.Request) -> str:
+    ip = request.client.ip
+    user_agent = request.headers.get('User-Agent', '')
+    return hashlib.sha256(f"{ip}{user_agent}".encode()).hexdigest()
+
+def generate_token(user_identifier: str) -> str:
+    payload = {
+        'exp': int(time.time()) + 3600,  # 1 hour expiration
+        'user_id': user_identifier
+    }
+    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
+
 def verify_token(token: str) -> bool:
     try:
         jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
@@ -61,6 +83,23 @@ def verify_token(token: str) -> bool:
     except jwt.PyJWTError:
         return False
 
+def check_rate_limit(ip, session):
+    global last_reset_time
+    with rate_limit_lock:
+        current_time = time.time()
+        if current_time - last_reset_time >= RESET_INTERVAL:
+            ip_usage.clear()
+            session_usage.clear()
+            last_reset_time = current_time
+        if ip_usage[ip] >= MAX_IP_USAGE:
+            return False, "IP rate limit exceeded. Please try again in 24 hours."
+        if session_usage[session] >= MAX_SESSION_USAGE:
+            return False, "Session rate limit exceeded. Please try again in 24 hours."
+        ip_usage[ip] += 1
+        session_usage[session] += 1
+        return True, ""
+
+
 def extract_and_format_code(text):
     # Extract code between triple backticks
     code_match = re.search(r'```python\s*(.*?)\s*```', text, re.DOTALL)
@@ -96,6 +135,11 @@ def generate_explanation(problem: str, solution: str, token: str) -> Dict[str, A
     if not verify_token(token):
         return {"error": "Invalid token"}
     
+    problem_solution_hash = hashlib.sha256(f"{problem}{solution}".encode()).hexdigest()
+    if token not in token_to_problem_solution or token_to_problem_solution[token] != problem_solution_hash:
+        return {"error": "No matching problem-solution pair found for this token"}
+
+    
     system_prompt = "You are a Python coding assistant specialized in explaining LeetCode problem solutions. Provide a clear and concise explanation of the given solution."
     full_prompt = f"""### Instruction:
 {system_prompt}
@@ -120,7 +164,15 @@ Here's the explanation of the solution:
     return {"explanation": generated_text}
 
 
-def generate_solution(instruction: str, token: str) -> Dict[str, Any]:
+def generate_solution(instruction: str, token: str, request: gr.Request) -> Dict[str, Any]:
+    ip = request.client.ip
+    session = request.client.session
+    user_identifier = generate_user_identifier(request)
+
+    is_allowed, message = check_rate_limit(ip, session)
+    if not is_allowed:
+        return {"error": message}
+
     if not verify_token(token):
         return {"error": "Invalid token"}
     
@@ -143,10 +195,19 @@ Here's the complete Python function implementation:
         generated_text += chunk["choices"][0]["text"]
     
     formatted_code = extract_and_format_code(generated_text)
-    user_data[token] = {"problem": instruction, "solution": formatted_code}
+    problem_solution_hash = hashlib.sha256(f"{instruction}{formatted_code}".encode()).hexdigest()
+    token_to_problem_solution[token] = problem_solution_hash
     return {"solution": formatted_code}
 
-def random_problem(token: str) -> Dict[str, Any]:
+def random_problem(token: str, request: gr.Request) -> Dict[str, Any]:
+    ip = request.client.ip
+    session = request.client.session
+    user_identifier = generate_user_identifier(request)
+
+    is_allowed, message = check_rate_limit(ip, session)
+    if not is_allowed:
+        return {"error": message}
+
     if not verify_token(token):
         return {"error": "Invalid token"}
     
@@ -159,22 +220,32 @@ def random_problem(token: str) -> Dict[str, Any]:
     
     return {"problem": problem}
 
-def explain_solution(token: str) -> Dict[str, Any]:
+def explain_solution(token: str, problem: str, solution: str, request: gr.Request) -> Dict[str, Any]:
+    ip = request.client.ip
+    session = request.client.session
+    user_identifier = generate_user_identifier(request)
+
+    is_allowed, message = check_rate_limit(ip, session)
+    if not is_allowed:
+        return {"error": message}
+
     if not verify_token(token):
         return {"error": "Invalid token"}
-    
-    if token not in user_data or not user_data[token].get("solution"):
-        return {"error": "No solution available to explain. Please generate a solution first."}
-    
-    problem = user_data[token]["problem"]
-    solution = user_data[token]["solution"]
-    
+
+    problem_solution_hash = hashlib.sha256(f"{problem}{solution}".encode()).hexdigest()
+    if token not in token_to_problem_solution or token_to_problem_solution[token] != problem_solution_hash:
+        return {"error": "No matching problem-solution pair found for this token"}
+
     return generate_explanation(problem, solution, token)
 
 # Create Gradio interfaces
 generate_interface = gr.Interface(
     fn=generate_solution,
-    inputs=[gr.Textbox(label="Problem Instruction"), gr.Textbox(label="JWT Token")],
+    inputs=[
+        gr.Textbox(label="Problem Instruction"),
+        gr.Textbox(label="JWT Token"),
+        gr.Request()
+    ],
     outputs=gr.JSON(),
     title="Generate Solution API",
     description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution."
@@ -182,7 +253,10 @@ generate_interface = gr.Interface(
 
 random_problem_interface = gr.Interface(
     fn=random_problem,
-    inputs=gr.Textbox(label="JWT Token"),
+    inputs=[
+        gr.Textbox(label="JWT Token"),
+        gr.Request()
+    ],
     outputs=gr.JSON(),
     title="Random Problem API",
     description="Provide a valid JWT token to get a random LeetCode problem."
@@ -190,10 +264,15 @@ random_problem_interface = gr.Interface(
 
 explain_interface = gr.Interface(
     fn=explain_solution,
-    inputs=gr.Textbox(label="JWT Token"),
+    inputs=[
+        gr.Textbox(label="JWT Token"),
+        gr.Textbox(label="Problem"),
+        gr.Textbox(label="Solution"),
+        gr.Request()
+    ],
     outputs=gr.JSON(),
     title="Explain Solution API",
-    description="Provide a valid JWT token to get an explanation of the last generated solution."
+    description="Provide a valid JWT token, problem, and solution to get an explanation of the solution."
 )
 
 demo = gr.TabbedInterface(