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

and whether the given user is subscribed

src/interface/desktop/chat.html

@@ -361,9 +361,22 @@
                                 if (newResponseText.getElementsByClassName("spinner").length > 0) {
                                     newResponseText.removeChild(loadingSpinner);
                                 }
-
+                                // Try to parse the chunk as a JSON object. It will be a JSON object if there is an error.
-                                newResponseText.innerHTML += chunk;
+                                if (chunk.startsWith("{") && chunk.endsWith("}")) {
-                                readStream();
+                                    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

src/khoj/configure.py

@@ -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:

src/khoj/interface/web/chat.html

@@ -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;
+                                // Try to parse the chunk as a JSON object. It will be a JSON object if there is an error.
-                                readStream();
+                                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

src/khoj/routers/api.py

@@ -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_minute=Depends(ApiUserRateLimiter(requests=10, subscribed_requests=60, window=60)),
-    rate_limiter_per_day=Depends(ApiUserRateLimiter(requests=500, window=60 * 60 * 24)),
+    rate_limiter_per_day=Depends(ApiUserRateLimiter(requests=10, subscribed_requests=600, window=60 * 60 * 24)),
 ) -> Response:
     user = request.user.object
 

src/khoj/routers/helpers.py

@@ -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())