diff --git a/src/khoj/interface/web/chat.html b/src/khoj/interface/web/chat.html
index 6e049718..1b2ccda0 100644
--- a/src/khoj/interface/web/chat.html
+++ b/src/khoj/interface/web/chat.html
@@ -932,6 +932,8 @@ To get started, just start typing below. You can also type / to see a list of co
websocketState.references = references;
} else if (chunk.type == "status") {
handleStreamResponse(websocketState.newResponseText, chunk.message, null, false);
+ } else if (chunk.type == "rate_limit") {
+ handleStreamResponse(websocketState.newResponseText, chunk.message, websocketState.loadingEllipsis, true);
} else {
rawResponse = chunk.response;
}
@@ -939,7 +941,7 @@ To get started, just start typing below. You can also type / to see a list of co
// If the chunk is not a JSON object, just display it as is
websocketState.rawResponse += chunk;
} finally {
- if (chunk.type != "status") {
+ if (chunk.type != "status" && chunk.type != "rate_limit") {
addMessageToChatBody(websocketState.rawResponse, websocketState.newResponseText, websocketState.references);
}
}
diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index c528356e..5258a7e8 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -5,7 +5,7 @@ from typing import Dict, Optional
from urllib.parse import unquote
from asgiref.sync import sync_to_async
-from fastapi import APIRouter, Depends, Request, WebSocket
+from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
from fastapi.requests import Request
from fastapi.responses import Response, StreamingResponse
from starlette.authentication import requires
@@ -292,14 +292,30 @@ async def websocket_endpoint(
connection_alive = False
logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
+ async def send_rate_limit_message(message: str):
+ nonlocal connection_alive
+ if not connection_alive:
+ return
+
+ status_packet = {
+ "type": "rate_limit",
+ "message": message,
+ "content-type": "application/json",
+ }
+ try:
+ await websocket.send_text(json.dumps(status_packet))
+ except ConnectionClosedOK:
+ connection_alive = False
+ logger.info(f"User {user} disconnected web socket. Emitting rest of responses to clear thread")
+
user: KhojUser = websocket.user.object
conversation = await ConversationAdapters.aget_conversation_by_user(
user, client_application=websocket.user.client_app, conversation_id=conversation_id
)
- hourly_limiter = ApiUserRateLimiter(requests=5, subscribed_requests=60, window=60, slug="chat_minute")
+ hourly_limiter = ApiUserRateLimiter(requests=1, subscribed_requests=60, window=60, slug="chat_minute")
- daily_limiter = ApiUserRateLimiter(requests=5, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
+ daily_limiter = ApiUserRateLimiter(requests=1, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
await is_ready_to_chat(user)
@@ -318,8 +334,12 @@ async def websocket_endpoint(
logger.debug(f"User {user} disconnected web socket")
break
- await sync_to_async(hourly_limiter)(websocket)
- await sync_to_async(daily_limiter)(websocket)
+ try:
+ await sync_to_async(hourly_limiter)(websocket)
+ await sync_to_async(daily_limiter)(websocket)
+ except HTTPException as e:
+ await send_rate_limit_message(e.detail)
+ break
conversation_commands = [get_conversation_command(query=q, any_references=True)]