Allow users to short-circuit LLM response to offline model

src/khoj/interface/web/chat.html

@@ -72,6 +72,9 @@
  function chat() {
  // Extract required fields for search from form
  let query = document.getElementById("chat-input").value.trim();
+ const abortButton = document.getElementById("abort-button");
+ let controller = new AbortController();
+ let signal = controller.signal;
  let resultsCount = localStorage.getItem("khojResultsCount") || 5;
  console.log(`Query: ${query}`);
 
@@ -111,18 +114,33 @@
  chatInput.classList.remove("option-enabled");
 
  // Call specified Khoj API which returns a streamed response of type text/plain
- fetch(url)
+ fetch(url, { signal })
  .then(response => {
  const reader = response.body.getReader();
  const decoder = new TextDecoder();
 
+ abortButton.addEventListener("click", () => {
+ controller.abort();
+ console.log("Download aborted");
+ if (newResponseText.getElementsByClassName("spinner").length > 0) {
+ newResponseText.removeChild(loadingSpinner);
+ }
+ newResponseText.innerHTML += "<b>Aborted.</b>";
+ document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+ document.getElementById("chat-input").removeAttribute("disabled");
+ abortButton.style.display = "none";
+ });
+
  function readStream() {
  reader.read().then(({ done, value }) => {
+
  if (done) {
  // Evaluate the contents of new_response_text.innerHTML after all the data has been streamed
  const currentHTML = newResponseText.innerHTML;
  newResponseText.innerHTML = formatHTMLMessage(currentHTML);
 
+ abortButton.style.display = "none";
+
  return;
  }
 
@@ -145,6 +163,7 @@
  // Display response from Khoj
  if (newResponseText.getElementsByClassName("spinner").length > 0) {
  newResponseText.removeChild(loadingSpinner);
+ abortButton.style.display = "block";
  }
 
  newResponseText.innerHTML += chunk;
@@ -153,10 +172,32 @@
 
  // Scroll to bottom of chat window as chat response is streamed
  document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+ }).catch(error => {
+ if (error.name === 'AbortError') {
+ abortButton.style.display = "none";
+
+ const cancelUrl = `/api/chat/cancel?client=web`;
+
+ console.log('Fetch aborted');
+ abortController = new AbortController();
+ signal = abortController.signal;
+ fetch(cancelUrl, { method: 'POST'})
+ .then(response => {
+ return response.json();
+ })
+ .then(data => {
+ console.log(data);
+ });
+ } else {
+ console.error('Fetch error:', error);
+ }
  });
  }
  readStream();
  document.getElementById("chat-input").removeAttribute("disabled");
+ })
+ .catch(error => {
+ console.error('Fetch error:', error);
  });
  }
 
@@ -297,6 +338,7 @@
  <div id="chat-tooltip" style="display: none;"></div>
  <textarea id="chat-input" class="option" oninput="onChatInput()" onkeyup=incrementalChat(event) autofocus="autofocus" placeholder="Type / to see a list of commands, or just type your questions and hit enter.">
  </textarea>
+ <button id="abort-button" class="option">Abort</button>
  </div>
  </body>
 
@@ -389,6 +431,9 @@
  border-bottom: 0;
  transform: rotate(-60deg);
  }
+ button#abort-button {
+ display: none;
+ }
  /* color chat bubble by you dark grey */
  .chat-message-text.you {
  color: #f8fafc;

src/khoj/processor/conversation/gpt4all/chat_model.py

@@ -178,6 +178,14 @@ def llm_thread(g, messages: List[ChatMessage], model: Any):
  raise e
 
  assert isinstance(model, GPT4All), "model should be of type GPT4All"
+
+ def callback_fn(token_id, token_string):
+ if state.should_cancel_chat:
+ logger.debug(f"Canceling chat as cancel_chat is set to {state.should_cancel_chat}")
+ state.should_cancel_chat = False
+ return False
+ return True
+
  user_message = messages[-1]
  system_message = messages[0]
  conversation_history = messages[1:-1]
@@ -196,7 +204,9 @@ def llm_thread(g, messages: List[ChatMessage], model: Any):
  prompted_message = templated_system_message + chat_history + templated_user_message
 
  state.chat_lock.acquire()
- response_iterator = model.generate(prompted_message, streaming=True, max_tokens=500, n_batch=512)
+ response_iterator = model.generate(
+ prompted_message, streaming=True, max_tokens=500, n_batch=512, callback=callback_fn
+ )
  try:
  for response in response_iterator:
  if any(stop_word in response.strip() for stop_word in stop_words):

src/khoj/processor/conversation/gpt4all/utils.py

@@ -13,7 +13,7 @@ def download_model(model_name: str):
 
  # Use GPU for Chat Model, if available
  try:
- model = GPT4All(model_name=model_name, device="gpu")
+ model = GPT4All(model_name=model_name, device="gpu", allow_download=True)
  logger.debug(f"Loaded {model_name} chat model to GPU.")
  except ValueError:
  model = GPT4All(model_name=model_name)

src/khoj/routers/api.py

@@ -648,6 +648,31 @@ def update(
  return {"status": "ok", "message": "khoj reloaded"}
 
 
+@api.post("/chat/cancel")
+def cancel_chat(
+ request: Request,
+ client: Optional[str] = None,
+ user_agent: Optional[str] = Header(None),
+ referer: Optional[str] = Header(None),
+ host: Optional[str] = Header(None),
+):
+ perform_chat_checks()
+
+ state.should_cancel_chat = True
+
+ update_telemetry_state(
+ request=request,
+ telemetry_type="api",
+ api="cancel_chat",
+ client=client,
+ user_agent=user_agent,
+ referer=referer,
+ host=host,
+ )
+
+ return Response(content=json.dumps({"status": "ok"}), media_type="application/json", status_code=200)
+
+
 @api.get("/chat/history")
 def chat_history(
  request: Request,

src/khoj/utils/state.py

@@ -31,6 +31,7 @@
 previous_query: str = None
 demo: bool = False
 khoj_version: str = None
+should_cancel_chat: bool = False
 
 if torch.cuda.is_available():
  # Use CUDA GPU