Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import re
 import logging
 import uuid
+import time
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any
 import asyncio
@@ -58,10 +59,79 @@ class TaskStatusResponse(BaseModel):
     # request_params: Optional[Dict[str, Any]] = None # Optionally return original params
 
 
+# Rate limiting dictionary
+class RateLimiter:
+    def __init__(self, max_requests: int, time_window: timedelta):
+        self.max_requests = max_requests
+        self.time_window = time_window
+        self.requests: Dict[str, list] = defaultdict(list)
+    
+    def _cleanup_old_requests(self, user_ip: str) -> None:
+        """Remove requests that are outside the time window."""
+        current_time = time.time()
+        self.requests[user_ip] = [
+            timestamp for timestamp in self.requests[user_ip]
+            if current_time - timestamp < self.time_window.total_seconds()
+        ]
+    
+    def is_rate_limited(self, user_ip: str) -> bool:
+        """Check if the user has exceeded their rate limit."""
+        self._cleanup_old_requests(user_ip)
+        
+        # Get current count after cleanup
+        current_count = len(self.requests[user_ip])
+        
+        # Add current request timestamp (incrementing the count)
+        current_time = time.time()
+        self.requests[user_ip].append(current_time)
+        
+        # Check if user has exceeded the maximum requests
+        return (current_count + 1) > self.max_requests
+    
+    def get_current_count(self, user_ip: str) -> int:
+        """Get the current request count for an IP."""
+        self._cleanup_old_requests(user_ip)
+        return len(self.requests[user_ip])
+
+
+# Initialize rate limiter with 100 requests per day
+rate_limiter = RateLimiter(
+    max_requests=12,
+    time_window=timedelta(days=1)
+)
+
+def get_user_ip(request: Request) -> str:
+    """Helper function to get user's IP address."""
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0]
+    return request.client.host
+
+
+class ApiRotator:
+    def __init__(self, apis):
+        self.apis = apis
+        self.last_successful_index = None
+
+    def get_prioritized_apis(self):
+        if self.last_successful_index is not None:
+            # Move the last successful API to the front
+            rotated_apis = (
+                [self.apis[self.last_successful_index]] + 
+                self.apis[:self.last_successful_index] + 
+                self.apis[self.last_successful_index+1:]
+            )
+            return rotated_apis
+        return self.apis
+
+    def update_last_successful(self, index):
+        self.last_successful_index = index
+
+
 # --- FastAPI App Initialization ---
 app = FastAPI(
     title="Dual Chat & Async Gemini API",
-    description="Provides direct chat via custom API and asynchronous tasks for Gemini API.",
+    description="Made by Cody from chrunos.com.",
     version="2.0.0"
 )
 
@@ -162,7 +232,17 @@ async def process_gemini_request_background(
 @app.post("/chat", response_class=StreamingResponse)
 async def direct_chat(payload: ChatPayload):
     logger.info(f"Direct chat request received. Temperature: {payload.temperature}, Message: '{payload.message[:50]}...'")
+    user_ip = get_user_ip(request)
     
+    if rate_limiter.is_rate_limited(user_ip):
+        current_count = rate_limiter.get_current_count(user_ip)
+        raise HTTPException(
+            status_code=429,
+            detail={
+                "error": "You have exceeded the maximum number of requests per day. Please try again tomorrow.",
+                "url": "https://t.me/chrunoss"
+            }
+        )
     custom_api_key_secret = os.getenv("CUSTOM_API_SECRET_KEY")
     custom_api_base_url = os.getenv("CUSTOM_API_BASE_URL", CUSTOM_API_BASE_URL_DEFAULT)
     custom_api_model = os.getenv("CUSTOM_API_MODEL", CUSTOM_API_MODEL_DEFAULT)