sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-11-23 15:38:55 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Add a more gracefull error message when the rate limit is exceeded

Browse source
This commit is contained in:
sabaimran 2024-04-08 15:20:54 +05:30
parent 11ce3e2268
commit 089e0d028b
2 changed files with 28 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/khoj/interface/web/chat.html
View file

@ -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);
}
}

30
src/khoj/routers/api_chat.py
View file

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