diff --git a/src/interface/desktop/chat.html b/src/interface/desktop/chat.html index ecd8ebf9..ac8d9589 100644 --- a/src/interface/desktop/chat.html +++ b/src/interface/desktop/chat.html @@ -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 diff --git a/src/khoj/configure.py b/src/khoj/configure.py index 2d21bd01..19c1fd81 100644 --- a/src/khoj/configure.py +++ b/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: diff --git a/src/khoj/interface/web/chat.html b/src/khoj/interface/web/chat.html index abab83ab..f67ab857 100644 --- a/src/khoj/interface/web/chat.html +++ b/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; - 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 diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py index d5b6ce0e..f1967e67 100644 --- a/src/khoj/routers/api.py +++ b/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_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 diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index c6fcb436..dbcef1bc 100644 --- a/src/khoj/routers/helpers.py +++ b/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())