Enforce API use limits depending on whether the server has billing enabled

and whether the given user is subscribed
This commit is contained in:
sabaimran 2023-11-24 21:55:16 -08:00
parent 9c868ee10b
commit e5b1350523
5 changed files with 45 additions and 11 deletions

View file

@ -361,9 +361,22 @@
if (newResponseText.getElementsByClassName("spinner").length > 0) {
newResponseText.removeChild(loadingSpinner);
}
newResponseText.innerHTML += chunk;
readStream();
// Try to parse the chunk as a JSON object. It will be a JSON object if there is an error.
if (chunk.startsWith("{") && chunk.endsWith("}")) {
try {
const responseAsJson = JSON.parse(chunk);
if (responseAsJson.detail) {
newResponseText.innerHTML += responseAsJson.detail;
}
} catch (error) {
// If the chunk is not a JSON object, just display it as is
newResponseText.innerHTML += chunk;
}
} else {
// If the chunk is not a JSON object, just display it as is
newResponseText.innerHTML += chunk;
readStream();
}
}
// Scroll to bottom of chat window as chat response is streamed

View file

@ -82,7 +82,8 @@ class UserAuthenticationBackend(AuthenticationBackend):
return AuthCredentials(["authenticated", "subscribed"]), AuthenticatedKhojUser(
user_with_token.user
)
return AuthCredentials(["authenticated"]), AuthenticatedKhojUser(user)
return AuthCredentials(["authenticated"]), AuthenticatedKhojUser(user)
return AuthCredentials(["authenticated", "subscribed"]), AuthenticatedKhojUser(user)
if len(request.headers.get("Authorization", "").split("Bearer ")) == 2:
# Get bearer token from header
bearer_token = request.headers["Authorization"].split("Bearer ")[1]
@ -101,7 +102,8 @@ class UserAuthenticationBackend(AuthenticationBackend):
return AuthCredentials(["authenticated", "subscribed"]), AuthenticatedKhojUser(
user_with_token.user
)
return AuthCredentials(["authenticated"]), AuthenticatedKhojUser(user_with_token.user)
return AuthCredentials(["authenticated"]), AuthenticatedKhojUser(user)
return AuthCredentials(["authenticated", "subscribed"]), AuthenticatedKhojUser(user)
if state.anonymous_mode:
user = await self.khojuser_manager.filter(username="default").prefetch_related("subscription").afirst()
if user:

View file

@ -403,8 +403,22 @@ To get started, just start typing below. You can also type / to see a list of co
newResponseText.removeChild(loadingSpinner);
}
newResponseText.innerHTML += chunk;
readStream();
// Try to parse the chunk as a JSON object. It will be a JSON object if there is an error.
if (chunk.startsWith("{") && chunk.endsWith("}")) {
try {
const responseAsJson = JSON.parse(chunk);
if (responseAsJson.detail) {
newResponseText.innerHTML += responseAsJson.detail;
}
} catch (error) {
// If the chunk is not a JSON object, just display it as is
newResponseText.innerHTML += chunk;
}
} else {
// If the chunk is not a JSON object, just display it as is
newResponseText.innerHTML += chunk;
readStream();
}
}
// Scroll to bottom of chat window as chat response is streamed

View file

@ -573,8 +573,8 @@ async def chat(
n: Optional[int] = 5,
d: Optional[float] = 0.18,
stream: Optional[bool] = False,
rate_limiter_per_minute=Depends(ApiUserRateLimiter(requests=30, window=60)),
rate_limiter_per_day=Depends(ApiUserRateLimiter(requests=500, window=60 * 60 * 24)),
rate_limiter_per_minute=Depends(ApiUserRateLimiter(requests=10, subscribed_requests=60, window=60)),
rate_limiter_per_day=Depends(ApiUserRateLimiter(requests=10, subscribed_requests=600, window=60 * 60 * 24)),
) -> Response:
user = request.user.object

View file

@ -11,6 +11,7 @@ from typing import Annotated, Any, Dict, Iterator, List, Optional, Tuple, Union
# External Packages
from fastapi import Depends, Header, HTTPException, Request
from starlette.authentication import has_required_scope
from khoj.database.adapters import ConversationAdapters
from khoj.database.models import KhojUser, Subscription
@ -270,13 +271,15 @@ def generate_chat_response(
class ApiUserRateLimiter:
def __init__(self, requests: int, window: int):
def __init__(self, requests: int, subscribed_requests: int, window: int):
self.requests = requests
self.subscribed_requests = subscribed_requests
self.window = window
self.cache: dict[str, list[float]] = defaultdict(list)
def __call__(self, request: Request):
user: KhojUser = request.user.object
subscribed = has_required_scope(request, ["subscribed"])
user_requests = self.cache[user.uuid]
# Remove requests outside of the time window
@ -285,8 +288,10 @@ class ApiUserRateLimiter:
user_requests.pop(0)
# Check if the user has exceeded the rate limit
if len(user_requests) >= self.requests:
if subscribed and len(user_requests) >= self.subscribed_requests:
raise HTTPException(status_code=429, detail="Too Many Requests")
if not subscribed and len(user_requests) >= self.requests:
raise HTTPException(status_code=429, detail="Too Many Requests. Subscribe to increase your rate limit.")
# Add the current request to the cache
user_requests.append(time())