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